Use uv

Co-authored-by: Awni Hannun <awni@apple.com>

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

@@ -2,9 +2,13 @@ name: 'Build CUDA wheel'
 description: 'Build CUDA wheel'
 
 inputs:
-  toolkit:
-    description: 'The CUDA toolkit'
+  arch:
+    description: 'Platform architecture tag'
     required: true
+    type: choice
+    options:
+      - x86_64
+      - aarch64
 
 runs:
   using: "composite"
@@ -12,9 +16,16 @@ runs:
     - name: Build package
       shell: bash
       env:
-        CMAKE_ARGS: -DMLX_BUILD_CUDA=ON -DCMAKE_CUDA_COMPILER=/usr/local/${{ inputs.toolkit }}/bin/nvcc
+        CMAKE_ARGS: -DMLX_BUILD_CUDA=ON
       run: |
         pip install auditwheel build patchelf setuptools
         python setup.py clean --all
         MLX_BUILD_STAGE=2 python -m build -w
-        bash python/scripts/repair_cuda.sh
+
+        auditwheel repair dist/mlx_cuda*.whl \
+          --plat manylinux_2_35_${{ inputs.arch }} \
+          --exclude libcublas* \
+          --exclude libcuda* \
+          --exclude libcudnn* \
+          --exclude libnccl* \
+          --exclude libnvrtc*

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

@@ -1,26 +0,0 @@
-name: 'Build and Test with CUDA'
-description: 'Build and test MLX with CUDA'
-
-inputs:
-  toolkit:
-    description: 'The CUDA toolkit'
-    required: true
-
-runs:
-  using: "composite"
-  steps:
-    - name: Install Python package
-      shell: bash
-      env:
-        DEBUG: 1
-        CMAKE_ARGS: -DMLX_BUILD_CUDA=ON -DCMAKE_COMPILE_WARNING_AS_ERROR=ON -DCMAKE_CUDA_COMPILER=/usr/local/${{ inputs.toolkit }}/bin/nvcc
-      run: pip install --no-build-isolation -e ".[dev]" -v
-
-    - name: Build CPP only
-      shell: bash
-      run: |
-        cmake . -B build \
-          -DMLX_BUILD_CUDA=ON \
-          -DCMAKE_CUDA_COMPILER=/usr/local/${{ inputs.toolkit }}/bin/nvcc \
-          -DCMAKE_BUILD_TYPE=DEBUG
-        cmake --build build -j $(nproc)

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

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

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

@@ -1,25 +1,35 @@
 name: 'Build and Test on Linux'
-description: 'Build and test MLX 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:
-        CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
         DEBUG: 1
-      run: pip install --no-build-isolation -e ".[dev]" -v
-
-    - name: Generate package stubs
-      shell: sh
+        CMAKE_ARGS: >-
+          -DCMAKE_COMPILE_WARNING_AS_ERROR=ON
+          -DMLX_BUILD_CUDA=${{ startsWith(inputs.toolkit, 'cuda') && 'ON' || 'OFF' }}
       run: |
-        pip install typing_extensions
-        python setup.py generate_stubs
-    
+        if ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }} ; then
+          # There is no GPU in arm64 runner, use a common arch.
+          CMAKE_ARGS="$CMAKE_ARGS -DMLX_CUDA_ARCHITECTURES=90a"
+          # Can not build tests and stubs when the built executables can not run.
+          CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF -DMLX_BUILD_PYTHON_STUBS=OFF"
+        fi
+        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: |
-        mkdir -p build && cd build
-        cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG
-        make -j $(nproc)
+        cmake . -B build -DCMAKE_BUILD_TYPE=Debug ${{ steps.python_build.outputs.CMAKE_ARGS }}
+        cmake --build build -j $(nproc)

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

@@ -17,6 +17,8 @@ runs:
   steps:
     - name: Build Python package
       shell: bash -l {0}
+      env:
+        MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
       run: |
         pip install build
         python setup.py clean --all
@@ -25,6 +27,8 @@ runs:
     - name: Build backend package
       if: ${{ inputs.build-backend }}
       shell: bash -l {0}
+      env:
+        MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
       run: |
         python setup.py clean --all
         MLX_BUILD_STAGE=2 python -m build -w

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

@@ -4,22 +4,14 @@ description: 'Build and test MLX on macOS'
 runs:
   using: "composite"
   steps:
-    - name: Install dependencies
+    - name: Build
       env:
         DEBUG: 1
         CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
       shell: bash -l {0}
       run: |
-        pip install --upgrade pip
-        pip install cmake setuptools nanobind==2.4.0
         pip install -e . -v
 
-    - name: Generate package stubs
-      shell: bash -l {0}
-      run: |
-        pip install typing_extensions
-        python setup.py generate_stubs
-
     - name: Install tests dependencies
       shell: bash -l {0}
       run: |

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

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

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

@@ -18,7 +18,18 @@ runs:
       shell: bash
       run: xcodebuild -showComponent MetalToolchain
 
-    - uses: conda-incubator/setup-miniconda@v3
-      with:
-        miniconda-version: "latest"
-        python-version: ${{ inputs.python-version }}
+    - name: Install Python
+      shell: sh
+      run: |
+        curl -LsSf https://astral.sh/uv/install.sh | sh
+        $HOME/.local/bin/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
+
+    - name: Install build dependencies
+      shell: sh
+      run: |
+        pip install --upgrade pip
+        pip install cmake setuptools typing_extensions

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

@@ -1,8 +1,8 @@
 name: 'Run Linux tests'
 
 inputs:
-  cpu-only:
-    description: 'Skip GPU tests'
+  has-gpu:
+    description: 'Run GPU tests'
     required: false
     default: false
 
@@ -17,7 +17,7 @@ runs:
         echo "::endgroup::"
 
     - name: Run distributed tests
-      if: ${{ inputs.cpu-only == 'true' }}
+      if: ${{ inputs.has-gpu == 'false' }}
       shell: bash
       run: |
         echo "::group::Distributed tests"
@@ -30,7 +30,7 @@ runs:
         echo "::endgroup::"
 
     - name: Run Python tests - CPU
-      if: ${{ inputs.cpu-only == 'true' }}
+      if: ${{ inputs.has-gpu == 'false' }}
       shell: bash
       env:
         DEVICE: cpu
@@ -40,7 +40,7 @@ runs:
         echo "::endgroup::"
 
     - name: Run Python tests - GPU
-      if: ${{ inputs.cpu-only == 'false' }}
+      if: ${{ inputs.has-gpu == 'true' }}
       shell: bash
       env:
         DEVICE: gpu
@@ -59,7 +59,7 @@ runs:
         echo "::endgroup::"
 
     - name: Run CPP tests - GPU
-      if: ${{ inputs.cpu-only == 'false' }}
+      if: ${{ inputs.has-gpu == 'true' }}
       shell: bash
       env:
         DEVICE: gpu

.github/scripts/build-sanitizer-tests.sh

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

.github/workflows/build_and_test.yml

@@ -13,74 +13,112 @@ permissions:
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: ${{ github.ref != 'refs/head/main' }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
 
 jobs:
   check_lint:
+    name: Check Lint
     runs-on: ubuntu-22.04
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@v6
       - uses: pre-commit/action@v3.0.1
 
   linux_build_and_test:
+    name: Linux (cpu, ${{ matrix.arch }})
     needs: check_lint
     strategy:
-      matrix:
-        runner:
-          - ubuntu-22.04
-          - ubuntu-22.04-arm
       fail-fast: false
-    runs-on: ${{ matrix.runner }}
+      matrix:
+        arch: ['x86_64', 'aarch64']
+    runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }}
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@v6
       - uses: ./.github/actions/setup-linux
       - uses: ./.github/actions/build-linux
       - uses: ./.github/actions/test-linux
+
+  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:
-          cpu-only: true
+          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"]
+        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@v5
+      - uses: actions/checkout@v6
       - uses: ./.github/actions/setup-macos
       - uses: ./.github/actions/build-macos
 
-  cuda_build_and_test:
-    if: github.repository == 'ml-explore/mlx'
-    strategy:
-      fail-fast: false
-      matrix:
-        toolkit: ['cuda-12.6', 'cuda-12.9']
-    runs-on: gpu-t4-4-core
-    needs: check_lint
-    steps:
-      - uses: actions/checkout@v5
-      - uses: ./.github/actions/setup-linux
-        with:
-          toolkit: ${{ matrix.toolkit }}
-      - uses: ./.github/actions/build-cuda
-        with:
-          toolkit: ${{ matrix.toolkit }}
-      - uses: ./.github/actions/test-linux
-
   build_documentation:
+    name: Build Documentation
     if: github.repository == 'ml-explore/mlx'
     runs-on: ubuntu-22.04
     needs: check_lint
     steps:
-      - uses: actions/checkout@v5
+      - 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 CPP Build (${{ matrix.arch }})
+    name: Linux Fedora (${{ matrix.arch }})
     needs: check_lint
     strategy:
       fail-fast: false
@@ -96,7 +134,7 @@ jobs:
       image: fedora:42
     steps:
       - name: Checkout code
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
 
       - name: CPP Build Test - No Release
         run: |

.github/workflows/documentation.yml

@@ -10,7 +10,7 @@ jobs:
   build:
     runs-on: ubuntu-22.04
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@v6
       - uses: ./.github/actions/build-docs
       
   deploy:

.github/workflows/nightly.yml

@@ -16,21 +16,21 @@ jobs:
         python_version: ["3.10", "3.14"]
     runs-on: ubuntu-22.04
     steps:
-      - uses: actions/checkout@v5
+      - 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@v5
+        uses: actions/upload-artifact@v6
         with:
           name: linux-wheels-${{ matrix.python_version }}
           path: wheelhouse/mlx-*.whl
           retention-days: 7
       - name: Upload mlx-cpu artifacts
         if: matrix.python_version == '3.10'
-        uses: actions/upload-artifact@v5
+        uses: actions/upload-artifact@v6
         with:
           name: mlx-cpu
           path: wheelhouse/mlx_cpu-*.whl
@@ -46,14 +46,12 @@ jobs:
           - ubuntu-22.04-arm
     runs-on: ${{ matrix.runner }}
     steps:
-      - uses: actions/checkout@v5
+      - 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
-        with:
-          cpu-only: true
 
   build_mac_release:
     if: github.repository == 'ml-explore/mlx'
@@ -62,11 +60,16 @@ jobs:
         python-version: ["3.10", "3.13"]
     runs-on: [self-hosted, macos]
     steps:
-      - uses: actions/checkout@v5
+      - 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:
@@ -82,7 +85,7 @@ jobs:
     if: github.repository == 'ml-explore/mlx'
     runs-on: ubuntu-22-large
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@v6
       - uses: ./.github/actions/setup-linux
         with:
           toolkit: 'cuda-12.9'
@@ -90,9 +93,10 @@ jobs:
         uses: ./.github/actions/build-cuda-release
         with:
           toolkit: 'cuda-12.9'
+          arch: 'x86_64'
       - name: Upload artifacts
-        uses: actions/upload-artifact@v5
+        uses: actions/upload-artifact@v6
         with:
           name: mlx-cuda
-          path: wheelhouse/mlx_cuda-*.whl
+          path: wheelhouse/mlx_cuda_*.whl
           retention-days: 7

.github/workflows/release.yml

@@ -6,29 +6,30 @@ on:
       - 'v*'
   workflow_dispatch:
     inputs:
+      publish:
+        description: 'Publish to PyPI (uncheck for dry run)'
+        required: false
+        type: boolean
+        default: true
       dev_release:
-        description: "Do a dev release or regular release"
-        required: true
-        default: "false"
+        description: 'Development release (DEV_RELEASE=1)'
+        required: false
+        type: boolean
+        default: false
 
 permissions:
   contents: read
 
 jobs:
-  setup:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Set publishing variables
-        run: echo "Publishing setup complete"
-
   build_documentation:
     if: github.repository == 'ml-explore/mlx'
-    runs-on: [self-hosted, macos]
+    runs-on: ubuntu-22.04
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@v6
       - uses: ./.github/actions/build-docs
-    
+
   deploy_documentation:
+    if: inputs.publish
     needs: build_documentation
     permissions:
       pages: write
@@ -51,30 +52,33 @@ jobs:
     runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }}
     env:
       PYPI_RELEASE: 1
-      DEV_RELEASE: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}
+      DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }}
     steps:
-      - uses: actions/checkout@v5
+      - 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' }}
+          build-backend: ${{ matrix.python_version == '3.10' }}
           arch: ${{ matrix.arch }}
       - name: Upload MLX artifacts
-        uses: actions/upload-artifact@v5
+        uses: actions/upload-artifact@v6
         with:
           overwrite: true
-          name: linux-wheels-${{ matrix.python_version }}
+          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@v5
+        uses: actions/upload-artifact@v6
         with:
           overwrite: true
-          name: mlx-cpu
+          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:
@@ -83,10 +87,9 @@ jobs:
     runs-on: [self-hosted, macos]
     env:
       PYPI_RELEASE: 1
-      DEV_RELEASE: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}
-
+      DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }}
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@v6
       - uses: ./.github/actions/setup-macos
         with:
           python-version: ${{ matrix.python-version }}
@@ -95,13 +98,8 @@ jobs:
         shell: bash -l {0}
         run: |
           pip install --upgrade pip
-          pip install cmake setuptools nanobind==2.4.0
+          pip install cmake setuptools typing_extensions
           pip install -e . -v
-      - name: Generate package stubs
-        shell: bash -l {0}
-        run: |
-          pip install typing_extensions
-          python setup.py generate_stubs
       - name: Build macOS 14 package
         uses: ./.github/actions/build-macos-release
         with:
@@ -113,85 +111,96 @@ jobs:
           macos-target: 15.0
           build-backend: ${{ matrix.python-version == '3.10' }}
       - name: Upload MLX artifacts
-        uses: actions/upload-artifact@v5
+        uses: actions/upload-artifact@v6
         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@v5
+        uses: actions/upload-artifact@v6
         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'
-    runs-on: ubuntu-22-large
+    strategy:
+      matrix:
+        arch: ['x86_64', 'aarch64']
+        toolkit: ['cuda-12.9', 'cuda-13.0']
+    runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22-large' || 'ubuntu-22-large-arm' }}
     env:
       PYPI_RELEASE: 1
-      DEV_RELEASE: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}
+      DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }}
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@v6
       - uses: ./.github/actions/setup-linux
         with:
-          toolkit: 'cuda-12.9'
+          toolkit: ${{ matrix.toolkit }}
+          use-ccache: false
       - name: Build Python package
         uses: ./.github/actions/build-cuda-release
         with:
-          toolkit: 'cuda-12.9'
+          arch: ${{ matrix.arch }}
       - name: Upload artifacts
-        uses: actions/upload-artifact@v5
+        uses: actions/upload-artifact@v6
         with:
           overwrite: true
-          name: mlx-cuda
-          path: wheelhouse/mlx_cuda-*.whl
+          name: mlx-${{ matrix.toolkit }}-${{ matrix.arch }}
+          path: wheelhouse/mlx_cuda_*.whl
+          if-no-files-found: error
 
   pypi-publish:
     name: Upload release to PyPI
     runs-on: ubuntu-latest
-    needs: [setup, build_linux_release, build_mac_release]
+    needs: [build_linux_release, build_mac_release]
     permissions:
       id-token: write
     environment:
-      name: pypi
+      name: ${{ inputs.publish && 'pypi' || '' }}
       url: https://pypi.org/p/mlx
     steps:
-      - uses: actions/download-artifact@v6
+      - uses: actions/download-artifact@v7
         with:
           pattern: linux-wheels-*
           merge-multiple: true
           path: dist
-      - uses: actions/download-artifact@v6
+      - uses: actions/download-artifact@v7
         with:
           pattern: mac-wheels-*
           merge-multiple: true
           path: dist
       - name: Display structure of downloaded files
-        run: ls -R dist
+        run: du -ah dist
       - name: Publish package distributions to PyPI
+        if: inputs.publish
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           repository-url: https://upload.pypi.org/legacy/
-  
+
   pypi-publish-cuda:
     name: Upload CUDA release to PyPI
     runs-on: ubuntu-latest
-    needs: [setup, build_cuda_release]
+    needs: [build_cuda_release]
     permissions:
       id-token: write
     environment:
-      name: pypi
+      name: ${{ inputs.publish && 'pypi' || '' }}
       url: https://pypi.org/p/mlx-cuda
     steps:
-      - uses: actions/download-artifact@v6
+      - uses: actions/download-artifact@v7
         with:
-          name: mlx-cuda
+          pattern: mlx-cuda-*
+          merge-multiple: true
           path: dist
       - name: Display structure of downloaded files
-        run: ls -R dist
+        run: du -ah dist
       - name: Publish package distributions to PyPI
+        if: inputs.publish
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           repository-url: https://upload.pypi.org/legacy/
@@ -199,20 +208,22 @@ jobs:
   pypi-publish-cpu:
     name: Upload CPU release to PyPI
     runs-on: ubuntu-latest
-    needs: [setup, build_linux_release]
+    needs: [build_linux_release]
     permissions:
       id-token: write
     environment:
-      name: pypi
+      name: ${{ inputs.publish && 'pypi' || '' }}
       url: https://pypi.org/p/mlx-cpu
     steps:
-      - uses: actions/download-artifact@v6
+      - uses: actions/download-artifact@v7
         with:
-          name: mlx-cpu
+          pattern: mlx-cpu-*
+          merge-multiple: true
           path: dist
       - name: Display structure of downloaded files
-        run: ls -R dist
+        run: du -ah dist
       - name: Publish package distributions to PyPI
+        if: inputs.publish
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           repository-url: https://upload.pypi.org/legacy/
@@ -220,20 +231,21 @@ jobs:
   pypi-publish-metal:
     name: Upload Metal release to PyPI
     runs-on: ubuntu-latest
-    needs: [setup, build_mac_release]
+    needs: [build_mac_release]
     permissions:
       id-token: write
     environment:
-      name: pypi
+      name: ${{ inputs.publish && 'pypi' || '' }}
       url: https://pypi.org/p/mlx-metal
     steps:
-      - uses: actions/download-artifact@v6
+      - uses: actions/download-artifact@v7
         with:
           name: mlx-metal
           path: dist
       - name: Display structure of downloaded files
-        run: ls -R dist
+        run: du -ah dist
       - name: Publish package distributions to PyPI
+        if: inputs.publish
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           repository-url: https://upload.pypi.org/legacy/

CMakeLists.txt

@@ -41,10 +41,14 @@ option(MLX_ENABLE_X64_MAC "Enable building for x64 macOS" OFF)
 option(MLX_BUILD_GGUF "Include support for GGUF format" ON)
 option(MLX_BUILD_SAFETENSORS "Include support for safetensors format" ON)
 option(MLX_BUILD_BLAS_FROM_SOURCE "Build OpenBLAS from source code" OFF)
+option(MLX_BUILD_PYTHON_STUBS "Build stub files for python bindings" ON)
 option(MLX_METAL_JIT "Use JIT compilation for Metal kernels" OFF)
 option(MLX_USE_CCACHE "Use CCache for compilation cache when available" ON)
 option(BUILD_SHARED_LIBS "Build mlx as a shared library" OFF)
 option(USE_SYSTEM_FMT "Use system's provided fmt library" OFF)
+option(USE_ASAN "Enable AddressSanitizer (ASan)" OFF)
+option(USE_UBSAN "Enable UndefinedBehaviorSanitizer (UBSan)" OFF)
+option(USE_TSAN "Enable ThreadSanitizer (TSan)" OFF)
 
 # --------------------- Processor tests -------------------------
 message(
@@ -81,6 +85,63 @@ if(MLX_USE_CCACHE)
   endif()
 endif()
 
+if(USE_ASAN AND USE_TSAN)
+  message(
+    FATAL_ERROR
+      "AddressSanitizer (ASan) and ThreadSanitizer (TSan) are mutually exclusive and cannot be enabled at the same time."
+  )
+endif()
+
+set(SANITIZER_COMPILE_FLAGS "")
+set(SANITIZER_LINK_FLAGS "")
+
+if(USE_ASAN)
+  if(WIN32 AND MSVC)
+    list(APPEND SANITIZER_COMPILE_FLAGS /fsanitize=address)
+    list(APPEND SANITIZER_LINK_FLAGS /fsanitize=address)
+  else()
+    list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=address)
+    list(APPEND SANITIZER_LINK_FLAGS -fsanitize=address)
+    if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+      list(APPEND SANITIZER_LINK_FLAGS -lpthread)
+    endif()
+  endif()
+endif()
+
+if(USE_UBSAN)
+  if(WIN32 AND MSVC)
+    if(CMAKE_CXX_COMPILER_ID STREQUAL "Clang")
+      list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=undefined)
+      list(APPEND SANITIZER_LINK_FLAGS -fsanitize=undefined)
+    else()
+      message(
+        WARNING
+          "UndefinedBehaviorSanitizer (UBSan) is not directly supported via a simple flag in MSVC."
+      )
+    endif()
+  else()
+    list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=undefined)
+    list(APPEND SANITIZER_LINK_FLAGS -fsanitize=undefined)
+  endif()
+endif()
+
+if(USE_TSAN)
+  if(WIN32 AND MSVC)
+    message(
+      FATAL_ERROR
+        "ThreadSanitizer (TSan) is not supported by the MSVC compiler. Please use Clang or GCC."
+    )
+  elseif(CMAKE_SYSTEM_NAME STREQUAL "Darwin")
+    message(FATAL_ERROR "ThreadSanitizer (TSan) is not supported on macOS.")
+  else()
+    list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=thread)
+    list(APPEND SANITIZER_LINK_FLAGS -fsanitize=thread)
+    if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+      list(APPEND SANITIZER_LINK_FLAGS -lpthread)
+    endif()
+  endif()
+endif()
+
 # ----------------------------- Lib -----------------------------
 
 include(FetchContent)
@@ -93,6 +154,8 @@ add_library(mlx)
 # ‘std::pair<float, float>’ when C++17 is enabled changed to match C++14 in GCC
 # 10.1
 target_compile_options(mlx PRIVATE -Wno-psabi)
+target_compile_options(mlx PUBLIC ${SANITIZER_COMPILE_FLAGS})
+target_link_options(mlx PUBLIC ${SANITIZER_LINK_FLAGS})
 
 if(MLX_BUILD_CUDA)
   enable_language(CUDA)
@@ -273,14 +336,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.10.2
+    GIT_SHALLOW TRUE
+    EXCLUDE_FROM_ALL)
+  FetchContent_MakeAvailable(nanobind)
   add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/python/src)
 endif()
 

benchmarks/python/comparative/bench_mlx.py

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

benchmarks/python/comparative/bench_torch.py

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

benchmarks/python/time_utils.py

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

docs/src/dev/extensions.rst

@@ -777,11 +777,11 @@ with the naive :meth:`simple_axpby` we first defined.
             mx.eval(z)
 
         # Timed run
-        s = time.time()
+        s = time.perf_counter()
         for i in range(100):
             z = f(x, y, alpha, beta)
             mx.eval(z)
-        e = time.time()
+        e = time.perf_counter()
         return 1000 * (e - s) / 100
 
     simple_time = bench(simple_axpby)

docs/src/dev/metal_logging.rst

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

docs/src/index.rst

@@ -89,5 +89,6 @@ are the CPU and GPU.
 
    dev/extensions
    dev/metal_debugger
+   dev/metal_logging
    dev/custom_metal_kernels
    dev/mlx_in_cpp

docs/src/install.rst

@@ -29,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.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)
 ^^^^^^^^^^^^^^^^
@@ -125,13 +128,6 @@ Run the tests with:
 
   python -m unittest discover python/tests
 
-Optional: Install stubs to enable auto completions and type checking from your
-IDE:
-
-.. code-block:: shell
-
-  python setup.py generate_stubs
-
 C++ API
 ^^^^^^^
 

docs/src/python/data_types.rst

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

docs/src/usage/compile.rst

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

docs/src/usage/distributed.rst

@@ -7,22 +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 three 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. It should be
-  faster for thunderbolt connections, but it also works over Ethernet.
-* `nccl <https://developer.nvidia.com/nccl>`_, for use in CUDA environments.
 
 The list of all currently supported operations and their documentation can be
 seen in the :ref:`API docs<distributed>`.
 
-.. 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
 ---------------
 
@@ -85,7 +92,7 @@ Selecting Backend
 ^^^^^^^^^^^^^^^^^
 
 You can select the backend you want to use when calling :func:`init` by passing
-one of ``{'any', 'ring', 'mpi', 'nccl'}``. When passing ``any``, MLX will try all
+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::
@@ -110,6 +117,8 @@ 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:
+
 Training Example
 ----------------
 
@@ -192,16 +201,273 @@ almost identical to the example above:
         loss = step(model, x, y)
         mx.eval(loss, model.parameters())
 
+.. _ring_section:
+
+Getting Started with Ring
+-------------------------
+
+The ring backend does not depend on any third party library so it is always
+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 are not supported in the ring backend.
+
+Defining a Ring
+^^^^^^^^^^^^^^^
+
+The easiest way to define and use a ring is via a JSON hostfile and the
+``mlx.launch`` :doc:`helper script <launching_distributed>`. For each node one
+defines a hostname to ssh into to run commands on this node and one or more IPs
+that this node will listen to for connections.
+
+For example the hostfile below defines a 4 node ring. ``hostname1`` will be
+rank 0, ``hostname2`` rank 1 etc.
+
+.. code:: json
+
+    [
+        {"ssh": "hostname1", "ips": ["123.123.123.1"]},
+        {"ssh": "hostname2", "ips": ["123.123.123.2"]},
+        {"ssh": "hostname3", "ips": ["123.123.123.3"]},
+        {"ssh": "hostname4", "ips": ["123.123.123.4"]}
+    ]
+
+Running ``mlx.launch --hostfile ring-4.json my_script.py`` will ssh into each
+node, run the script which will listen for connections in each of the provided
+IPs. Specifically, ``hostname1`` will connect to ``123.123.123.2`` and accept a
+connection from ``123.123.123.4`` and so on and so forth.
+
+Thunderbolt Ring
+^^^^^^^^^^^^^^^^
+
+Although the ring backend can have benefits over MPI even for Ethernet, its
+main purpose is to use Thunderbolt rings for higher bandwidth communication.
+Setting up such thunderbolt rings can be done manually, but is a relatively
+tedious process. To simplify this, we provide the utility ``mlx.distributed_config``.
+
+To use ``mlx.distributed_config`` your computers need to be accessible by ssh via
+Ethernet or Wi-Fi. Subsequently, connect them via thunderbolt cables and then call the
+utility as follows:
+
+.. code:: shell
+
+   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.
+
+If you want to go through the process manually, the steps are as follows:
+
+* Disable the thunderbolt bridge interface
+* For the cable connecting rank ``i`` to rank ``i + 1`` find the interfaces
+  corresponding to that cable in nodes ``i`` and ``i + 1``.
+* Set up a unique subnetwork connecting the two nodes for the corresponding
+  interfaces. For instance if the cable corresponds to ``en2`` on node ``i``
+  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 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.
+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:
@@ -269,78 +535,116 @@ 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.
 
-Getting Started with Ring
--------------------------
+.. _no_mlx_launch:
 
-The ring backend does not depend on any third party library so it is always
-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.
+Distributed Without ``mlx.launch``
+----------------------------------
 
-Defining a Ring
-^^^^^^^^^^^^^^^
+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.
 
-The easiest way to define and use a ring is via a JSON hostfile and the
-``mlx.launch`` :doc:`helper script <launching_distributed>`. For each node one
-defines a hostname to ssh into to run commands on this node and one or more IPs
-that this node will listen to for connections.
+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.
 
-For example the hostfile below defines a 4 node ring. ``hostname1`` will be
-rank 0, ``hostname2`` rank 1 etc.
+Below we list the environment variables required to use each backend.
 
-.. code:: json
+Ring
+^^^^^^
 
-    [
-        {"ssh": "hostname1", "ips": ["123.123.123.1"]},
-        {"ssh": "hostname2", "ips": ["123.123.123.2"]},
-        {"ssh": "hostname3", "ips": ["123.123.123.3"]},
-        {"ssh": "hostname4", "ips": ["123.123.123.4"]}
-    ]
+**MLX_RANK** should contain a single 0-based integer that defines the rank of
+the process.
 
-Running ``mlx.launch --hostfile ring-4.json my_script.py`` will ssh into each
-node, run the script which will listen for connections in each of the provided
-IPs. Specifically, ``hostname1`` will connect to ``123.123.123.2`` and accept a
-connection from ``123.123.123.4`` and so on and so forth.
+**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:
 
-Thunderbolt Ring
-^^^^^^^^^^^^^^^^
+.. code-block:: json
 
-Although the ring backend can have benefits over MPI even for Ethernet, its
-main purpose is to use Thunderbolt rings for higher bandwidth communication.
-Setting up such thunderbolt rings can be done manually, but is a relatively
-tedious process. To simplify this, we provide the utility ``mlx.distributed_config``.
+   [
+     ["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"]
+   ]
 
-To use ``mlx.distributed_config`` your computers need to be accessible by ssh via
-Ethernet or Wi-Fi. Subsequently, connect them via thunderbolt cables and then call the
-utility as follows:
+**MLX_RING_VERBOSE** is optional and if set to 1 it enables some more logging
+from the distributed backend.
 
-.. code:: shell
+JACCL
+^^^^^
 
-   mlx.distributed_config --verbose --hosts host1,host2,host3,host4
+**MLX_RANK** should contain a single 0-based integer that defines the rank of
+the process.
 
-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.
+**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.
 
-To validate your connection without configuring anything
-``mlx.distributed_config`` can also plot the ring using DOT format.
+**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:: shell
+.. code-block:: json
 
-   mlx.distributed_config --verbose --hosts host1,host2,host3,host4 --dot >ring.dot
-   dot -Tpng ring.dot >ring.png
-   open ring.png
+   [
+      [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]
+   ]
 
-If you want to go through the process manually, the steps are as follows:
 
-* Disable the thunderbolt bridge interface
-* For the cable connecting rank ``i`` to rank ``i + 1`` find the interfaces
-  corresponding to that cable in nodes ``i`` and ``i + 1``.
-* Set up a unique subnetwork connecting the two nodes for the corresponding
-  interfaces. For instance if the cable corresponds to ``en2`` on node ``i``
-  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.
+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.

docs/src/usage/indexing.rst

@@ -179,14 +179,14 @@ assignments, ``updates`` must provide at least as many elements as there are
 
 Boolean masks follow NumPy semantics:
 
-- The mask shape must match the shape of the axes it indexes exactly. No mask
-  broadcasting occurs.
+- 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.randn(10, 10) > 0.0] = 0  # valid: mask covers axes 0 and 1
+   >>> 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``.

docs/src/usage/launching_distributed.rst

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

examples/extensions/pyproject.toml

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

examples/extensions/requirements.txt

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

examples/python/linear_regression.py

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

examples/python/logistic_regression.py

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

examples/python/qqmm.py

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

mlx/CMakeLists.txt

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

mlx/allocator.cpp

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

mlx/allocator.h

@@ -28,16 +28,16 @@ class Buffer {
   };
 };
 
-Buffer malloc(size_t size);
-
-void free(Buffer buffer);
-
 class 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;
@@ -49,4 +49,25 @@ class Allocator {
 
 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

mlx/array.cpp

@@ -82,6 +82,28 @@ array::array(std::initializer_list<int> data, Dtype 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)) {

mlx/array.h

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

mlx/backend/common/compiled.cpp

@@ -130,7 +130,7 @@ void compiled_allocate_outputs(
       // - Donatable
       // - Not a constant
       if (in.itemsize() == outputs[o].itemsize() && !is_scalar(in) &&
-          in.is_donatable() && is_constant(i)) {
+          in.is_donatable() && !is_constant(i)) {
         outputs[o++].copy_shared_buffer(in);
       }
       // Get representative input flags to properly set non-donated outputs
@@ -158,7 +158,7 @@ void compiled_allocate_outputs(
       // - Not a constant
       if (in.flags().row_contiguous && in.size() == outputs[o].size() &&
           in.itemsize() == outputs[o].itemsize() && in.is_donatable() &&
-          is_constant(i)) {
+          !is_constant(i)) {
         outputs[o].copy_shared_buffer(
             in, outputs[o].strides(), in.flags(), in.data_size());
         o++;

mlx/backend/cpu/eig.cpp

@@ -12,6 +12,167 @@ namespace mlx::core {
 
 namespace {
 
+template <typename T>
+complex64_t to_complex(T r, T i) {
+  return {static_cast<float>(r), static_cast<float>(i)};
+}
+
+template <typename T, class Enable = void>
+struct EigWork {};
+
+template <typename T>
+struct EigWork<
+    T,
+    typename std::enable_if<std::is_floating_point<T>::value>::type> {
+  using O = complex64_t;
+
+  char jobl;
+  char jobr;
+  int N;
+  int lwork;
+  int info;
+  std::vector<array::Data> buffers;
+
+  EigWork(char jobl_, char jobr_, int N_, bool compute_eigenvectors)
+      : jobl(jobl_), jobr(jobr_), N(N_), lwork(-1) {
+    T work;
+    int n_vecs_l = compute_eigenvectors ? N_ : 1;
+    int n_vecs_r = 1;
+    geev<T>(
+        &jobl,
+        &jobr,
+        &N,
+        nullptr,
+        &N,
+        nullptr,
+        nullptr,
+        nullptr,
+        &n_vecs_l,
+        nullptr,
+        &n_vecs_r,
+        &work,
+        &lwork,
+        &info);
+    lwork = static_cast<int>(work);
+
+    buffers.emplace_back(allocator::malloc(sizeof(T) * N * 2));
+    if (compute_eigenvectors) {
+      buffers.emplace_back(allocator::malloc(sizeof(T) * N * N * 2));
+    }
+    buffers.emplace_back(allocator::malloc(sizeof(T) * lwork));
+  }
+
+  void run(T* a, O* values, O* vectors) {
+    auto eig_tmp = static_cast<T*>(buffers[0].buffer.raw_ptr());
+    T* vec_tmp = nullptr;
+    if (vectors) {
+      vec_tmp = static_cast<T*>(buffers[1].buffer.raw_ptr());
+    }
+    auto work = static_cast<T*>(buffers.back().buffer.raw_ptr());
+
+    int n_vecs_l = vectors ? N : 1;
+    int n_vecs_r = 1;
+    geev<T>(
+        &jobl,
+        &jobr,
+        &N,
+        a,
+        &N,
+        eig_tmp,
+        eig_tmp + N,
+        vectors ? vec_tmp : nullptr,
+        &n_vecs_l,
+        nullptr,
+        &n_vecs_r,
+        work,
+        &lwork,
+        &info);
+
+    for (int i = 0; i < N; ++i) {
+      values[i] = to_complex(eig_tmp[i], eig_tmp[N + i]);
+    }
+
+    if (vectors) {
+      for (int i = 0; i < N; ++i) {
+        if (values[i].imag() != 0) {
+          for (int j = 0; j < N; ++j) {
+            vectors[i * N + j] =
+                to_complex(vec_tmp[i * N + j], -vec_tmp[(i + 1) * N + j]);
+            vectors[(i + 1) * N + j] =
+                to_complex(vec_tmp[i * N + j], vec_tmp[(i + 1) * N + j]);
+          }
+          i += 1;
+        } else {
+          for (int j = 0; j < N; ++j) {
+            vectors[i * N + j] = to_complex(vec_tmp[i * N + j], T(0.0));
+          }
+        }
+      }
+    }
+  }
+};
+
+template <>
+struct EigWork<std::complex<float>> {
+  using T = std::complex<float>;
+  using R = float;
+  using O = T;
+
+  char jobl;
+  char jobr;
+  int N;
+  int lwork;
+  int lrwork;
+  int info;
+  std::vector<array::Data> buffers;
+
+  EigWork(char jobl_, char jobr_, int N_, bool compute_eigenvectors)
+      : jobl(jobl_), jobr(jobr_), N(N_), lwork(-1), lrwork(2 * N_) {
+    T work;
+    R rwork;
+    int n_vecs_l = compute_eigenvectors ? N_ : 1;
+    int n_vecs_r = 1;
+    geev<T>(
+        &jobl,
+        &jobr,
+        &N,
+        nullptr,
+        &N,
+        nullptr,
+        nullptr,
+        &n_vecs_l,
+        nullptr,
+        &n_vecs_r,
+        &work,
+        &lwork,
+        &rwork,
+        &info);
+    lwork = static_cast<int>(work.real());
+    buffers.emplace_back(allocator::malloc(sizeof(T) * lwork));
+    buffers.emplace_back(allocator::malloc(sizeof(R) * lrwork));
+  }
+
+  void run(T* a, T* values, T* vectors) {
+    int n_vecs_l = vectors ? N : 1;
+    int n_vecs_r = 1;
+    geev<T>(
+        &jobl,
+        &jobr,
+        &N,
+        a,
+        &N,
+        values,
+        vectors,
+        &n_vecs_l,
+        nullptr,
+        &n_vecs_r,
+        static_cast<T*>(buffers[0].buffer.raw_ptr()),
+        &lwork,
+        static_cast<R*>(buffers[1].buffer.raw_ptr()),
+        &info);
+  }
+};
+
 template <typename T>
 void eig_impl(
     array& a,
@@ -19,101 +180,39 @@ void eig_impl(
     array& values,
     bool compute_eigenvectors,
     Stream stream) {
-  using OT = std::complex<T>;
   auto a_ptr = a.data<T>();
-  auto eig_ptr = values.data<OT>();
+  auto val_ptr = values.data<complex64_t>();
 
   auto& encoder = cpu::get_command_encoder(stream);
   encoder.set_input_array(a);
   encoder.set_output_array(values);
-  OT* vec_ptr = nullptr;
+  complex64_t* vec_ptr = nullptr;
   if (compute_eigenvectors) {
     encoder.set_output_array(vectors);
-    vec_ptr = vectors.data<OT>();
+    vec_ptr = vectors.data<complex64_t>();
   }
   encoder.dispatch([a_ptr,
+                    val_ptr,
                     vec_ptr,
-                    eig_ptr,
                     compute_eigenvectors,
                     N = vectors.shape(-1),
                     size = vectors.size()]() mutable {
-    // Work query
     char jobr = 'N';
     char jobl = compute_eigenvectors ? 'V' : 'N';
-    int n_vecs_r = 1;
-    int n_vecs_l = compute_eigenvectors ? N : 1;
-    int lwork = -1;
-    int info;
-    {
-      T work;
-      geev<T>(
-          &jobl,
-          &jobr,
-          &N,
-          nullptr,
-          &N,
-          nullptr,
-          nullptr,
-          nullptr,
-          &n_vecs_l,
-          nullptr,
-          &n_vecs_r,
-          &work,
-          &lwork,
-          &info);
-      lwork = static_cast<int>(work);
-    }
 
-    auto eig_tmp_data = array::Data{allocator::malloc(sizeof(T) * N * 2)};
-    auto vec_tmp_data =
-        array::Data{allocator::malloc(vec_ptr ? sizeof(T) * N * N * 2 : 0)};
-    auto eig_tmp = static_cast<T*>(eig_tmp_data.buffer.raw_ptr());
-    auto vec_tmp = static_cast<T*>(vec_tmp_data.buffer.raw_ptr());
-    auto work_buf = array::Data{allocator::malloc(sizeof(T) * lwork)};
+    EigWork<T> work(jobl, jobr, N, compute_eigenvectors);
+
     for (size_t i = 0; i < size / (N * N); ++i) {
-      geev<T>(
-          &jobl,
-          &jobr,
-          &N,
-          a_ptr,
-          &N,
-          eig_tmp,
-          eig_tmp + N,
-          vec_tmp,
-          &n_vecs_l,
-          nullptr,
-          &n_vecs_r,
-          static_cast<T*>(work_buf.buffer.raw_ptr()),
-          &lwork,
-          &info);
-      for (int i = 0; i < N; ++i) {
-        eig_ptr[i] = {eig_tmp[i], eig_tmp[N + i]};
-      }
+      work.run(a_ptr, val_ptr, vec_ptr);
+      a_ptr += N * N;
+      val_ptr += N;
       if (vec_ptr) {
-        for (int i = 0; i < N; ++i) {
-          if (eig_ptr[i].imag() != 0) {
-            // This vector and the next are a pair
-            for (int j = 0; j < N; ++j) {
-              vec_ptr[i * N + j] = {
-                  vec_tmp[i * N + j], -vec_tmp[(i + 1) * N + j]};
-              vec_ptr[(i + 1) * N + j] = {
-                  vec_tmp[i * N + j], vec_tmp[(i + 1) * N + j]};
-            }
-            i += 1;
-          } else {
-            for (int j = 0; j < N; ++j) {
-              vec_ptr[i * N + j] = {vec_tmp[i * N + j], 0};
-            }
-          }
-        }
         vec_ptr += N * N;
       }
-      a_ptr += N * N;
-      eig_ptr += N;
-      if (info != 0) {
+      if (work.info != 0) {
         std::stringstream msg;
         msg << "[Eig::eval_cpu] Eigenvalue decomposition failed with error code "
-            << info;
+            << work.info;
         throw std::runtime_error(msg.str());
       }
     }
@@ -165,8 +264,17 @@ void Eig::eval_cpu(
     case float32:
       eig_impl<float>(a_copy, vectors, values, compute_eigenvectors_, stream());
       break;
+    case float64:
+      eig_impl<double>(
+          a_copy, vectors, values, compute_eigenvectors_, stream());
+      break;
+    case complex64:
+      eig_impl<std::complex<float>>(
+          a_copy, vectors, values, compute_eigenvectors_, stream());
+      break;
     default:
-      throw std::runtime_error("[Eig::eval_cpu] only supports float32.");
+      throw std::runtime_error(
+          "[Eig::eval_cpu] only supports float32, float64, or complex64.");
   }
 }
 

mlx/backend/cpu/lapack.h

@@ -45,9 +45,7 @@
 INSTANTIATE_LAPACK_REAL(geqrf)
 INSTANTIATE_LAPACK_REAL(orgqr)
 INSTANTIATE_LAPACK_REAL(syevd)
-INSTANTIATE_LAPACK_REAL(geev)
 INSTANTIATE_LAPACK_REAL(potrf)
-INSTANTIATE_LAPACK_REAL(gesdd)
 INSTANTIATE_LAPACK_REAL(getrf)
 INSTANTIATE_LAPACK_REAL(getri)
 INSTANTIATE_LAPACK_REAL(trtri)
@@ -63,3 +61,20 @@ INSTANTIATE_LAPACK_REAL(trtri)
   }
 
 INSTANTIATE_LAPACK_COMPLEX(heevd)
+
+#define INSTANTIATE_LAPACK_ALL(FUNC)                                \
+  template <typename T, typename... Args>                           \
+  void FUNC(Args... args) {                                         \
+    if constexpr (std::is_same_v<T, float>) {                       \
+      MLX_LAPACK_FUNC(s##FUNC)(std::forward<Args>(args)...);        \
+    } else if constexpr (std::is_same_v<T, double>) {               \
+      MLX_LAPACK_FUNC(d##FUNC)(std::forward<Args>(args)...);        \
+    } else if constexpr (std::is_same_v<T, std::complex<float>>) {  \
+      MLX_LAPACK_FUNC(c##FUNC)(std::forward<Args>(args)...);        \
+    } else if constexpr (std::is_same_v<T, std::complex<double>>) { \
+      MLX_LAPACK_FUNC(z##FUNC)(std::forward<Args>(args)...);        \
+    }                                                               \
+  }
+
+INSTANTIATE_LAPACK_ALL(geev)
+INSTANTIATE_LAPACK_ALL(gesdd)

mlx/backend/cpu/primitives.cpp

@@ -291,6 +291,17 @@ void RandomBits::eval_cpu(const std::vector<array>& inputs, array& out) {
                     num_keys,
                     kshape = keys.shape(),
                     kstrides = keys.strides()]() mutable {
+    auto copy_remaining = [&](char* cptr, size_t loc, uint32_t v) {
+      if (4 * loc + 4 <= bytes_per_key) {
+        reinterpret_cast<uint32_t*>(cptr)[loc] = v;
+      } else {
+        std::copy(
+            reinterpret_cast<char*>(&v),
+            reinterpret_cast<char*>(&v) + bytes_per_key - 4 * loc,
+            cptr + 4 * loc);
+      }
+    };
+
     size_t out_skip = (bytes_per_key + 4 - 1) / 4;
     auto half_size = out_skip / 2;
     bool even = out_skip % 2 == 0;
@@ -310,18 +321,12 @@ void RandomBits::eval_cpu(const std::vector<array>& inputs, array& out) {
       if (count.first < half_size) {
         auto rb = random::threefry2x32_hash(key, count);
         ptr[count.first++] = rb.first;
-        if (bytes_per_key % 4 > 0) {
-          std::copy(
-              reinterpret_cast<char*>(&rb.second),
-              reinterpret_cast<char*>(&rb.second) + bytes_per_key % 4,
-              cptr + 4 * count.second);
-        } else {
-          ptr[count.second] = rb.second;
-        }
+        copy_remaining(cptr, count.second, rb.second);
       }
       if (!even) {
         count.second = 0;
-        ptr[half_size] = random::threefry2x32_hash(key, count).first;
+        copy_remaining(
+            cptr, half_size, random::threefry2x32_hash(key, count).first);
       }
     }
   });

mlx/backend/cpu/quantized.cpp

@@ -14,7 +14,7 @@ namespace mlx::core {
 
 namespace {
 
-const static float MXFP4_LUT[16] = {
+const static float FP4_LUT[16] = {
     +0.0f,
     +0.5f,
     +1.0f,
@@ -32,15 +32,19 @@ const static float MXFP4_LUT[16] = {
     -4.0f,
     -6.0f};
 
-template <typename T>
+template <typename T, int group_size>
 static inline T dequantize_scale(uint8_t s) {
-  using FOrI = union {
-    bfloat16_t f;
-    uint16_t i;
-  };
-  FOrI out;
-  out.i = (s == 0 ? 0x40 : (static_cast<uint16_t>(s) << 7));
-  return static_cast<T>(out.f);
+  if constexpr (group_size == 16) {
+    return static_cast<T>(detail::FromFP8{}(s));
+  } else {
+    using FOrI = union {
+      bfloat16_t f;
+      uint16_t i;
+    };
+    FOrI out;
+    out.i = (s == 0 ? 0x40 : (static_cast<uint16_t>(s) << 7));
+    return static_cast<T>(out.f);
+  }
 }
 
 inline constexpr short get_pack_factor(int bits, int wsize = 8) {
@@ -437,8 +441,8 @@ void _qmm_dispatch(
   }
 }
 
-template <typename T>
-void mxfp4_qmm(
+template <typename T, int group_size, int bits>
+void fp_qmm(
     T* result,
     const T* x,
     const uint32_t* w,
@@ -446,8 +450,7 @@ void mxfp4_qmm(
     int M,
     int N,
     int K) {
-  constexpr int group_size = 32;
-  constexpr int pack_factor = get_pack_factor(4, 8);
+  constexpr int pack_factor = get_pack_factor(bits, 8);
   constexpr int packs_in_group = group_size / pack_factor;
 
   for (int m = 0; m < M; m++) {
@@ -461,25 +464,27 @@ void mxfp4_qmm(
       T xi = *x++;
 
       for (int n = 0; n < N; n += group_size) {
-        T scale = dequantize_scale<T>(*scales_local++);
+        T scale = dequantize_scale<T, group_size>(*scales_local++);
         for (int ng = 0; ng < packs_in_group; ng++) {
-          uint8_t wi = *w_local++;
-#pragma clang loop unroll(full)
-          for (int p = 0; p < pack_factor; p++) {
+          if constexpr (bits == 4) {
             (*result_local++) +=
-                xi * scale * static_cast<T>(MXFP4_LUT[wi & 0xf]);
-            wi >>= 4;
+                xi * scale * static_cast<T>(FP4_LUT[w_local[0] & 0xf]);
+            (*result_local++) +=
+                xi * scale * static_cast<T>(FP4_LUT[(w_local[0] >> 4) & 0xf]);
+          } else {
+            (*result_local++) +=
+                xi * scale * static_cast<T>(detail::FromFP8{}(w_local[0]));
           }
+          w_local++;
         }
       }
     }
-
     result += N;
   }
 }
 
-template <typename T>
-void mxfp4_qmm_t(
+template <typename T, int group_size, int bits>
+void fp_qmm_t(
     T* result,
     const T* x,
     const uint32_t* w,
@@ -487,8 +492,7 @@ void mxfp4_qmm_t(
     int M,
     int N,
     int K) {
-  constexpr int group_size = 32;
-  constexpr int pack_factor = get_pack_factor(4, 8);
+  constexpr int pack_factor = get_pack_factor(bits, 8);
   constexpr int packs_in_group = group_size / pack_factor;
 
   for (int m = 0; m < M; m++) {
@@ -499,16 +503,19 @@ void mxfp4_qmm_t(
       const T* x_local = x;
       T sum = 0;
       for (int k = 0; k < K; k += group_size) {
-        T scale = dequantize_scale<T>(*scales_local++);
+        T scale = dequantize_scale<T, group_size>(*scales_local++);
 
         T gsum = 0;
         for (int kw = 0; kw < packs_in_group; kw++) {
-          uint8_t wi = *w_local++;
-#pragma clang loop unroll(full)
-          for (int p = 0; p < pack_factor; p++) {
-            gsum += (*x_local++) * static_cast<T>(MXFP4_LUT[wi & 0xf]);
-            wi >>= 4;
+          if constexpr (bits == 4) {
+            gsum += (*x_local++) * static_cast<T>(FP4_LUT[w_local[0] & 0xf]);
+            gsum +=
+                (*x_local++) * static_cast<T>(FP4_LUT[(w_local[0] >> 4) & 0xf]);
+          } else {
+            gsum +=
+                (*x_local++) * static_cast<T>(detail::FromFP8{}(w_local[0]));
           }
+          w_local++;
         }
         sum += scale * gsum;
       }
@@ -520,9 +527,9 @@ void mxfp4_qmm_t(
   }
 }
 
-template <int S>
-simd::Simd<float, S> mxfp4_extract_bits_simd(const uint32_t* w) {
-  if constexpr (S == 8) {
+template <int S, int bits>
+simd::Simd<float, S> fp_extract_bits_simd(const uint32_t* w) {
+  if constexpr (S == 8 && bits == 4) {
     constexpr std::array<uint32_t, 8> shifts_ = {{0, 4, 8, 12, 16, 20, 24, 28}};
     auto shifts(*(simd::Simd<uint32_t, S>*)&shifts_);
     auto wi = simd::Simd<uint32_t, S>(*w);
@@ -530,17 +537,20 @@ simd::Simd<float, S> mxfp4_extract_bits_simd(const uint32_t* w) {
     wi = wi & 0xf;
     simd::Simd<float, S> w_out;
     for (int i = 0; i < S; ++i) {
-      w_out[i] = MXFP4_LUT[wi[i]];
+      w_out[i] = FP4_LUT[wi[i]];
     }
     return w_out;
+  } else if constexpr (S == 8 && bits == 8) {
+    auto w_out = simd::load<uint8_t, S>(reinterpret_cast<const uint8_t*>(w));
+    return detail::FromFP8{}(w_out);
   } else {
     // Appease compiler.. but should never get here
     throw std::runtime_error("Unsupported combination for simd qmm.");
   }
 }
 
-template <typename T>
-void mxfp4_qmm_t_simd(
+template <typename T, int group_size, int bits>
+void fp_qmm_t_simd(
     T* result,
     const T* x,
     const uint32_t* w,
@@ -548,8 +558,7 @@ void mxfp4_qmm_t_simd(
     int M,
     int N,
     int K) {
-  constexpr int group_size = 32;
-  constexpr int pack_factor = 32 / 4;
+  constexpr int pack_factor = get_pack_factor(bits, 32);
   constexpr int packs_in_group = group_size / pack_factor;
   constexpr int S = simd::max_size<T>;
   static_assert(
@@ -564,12 +573,12 @@ void mxfp4_qmm_t_simd(
       simd::Simd<float, S> acc(0);
       auto x_local = x;
       for (int k = 0; k < K; k += group_size) {
-        T scale = dequantize_scale<T>(*scales_local++);
+        T scale = dequantize_scale<T, group_size>(*scales_local++);
 
         simd::Simd<float, S> g_acc(0);
         for (int kw = 0; kw < packs_in_group; kw += packs_per_simd) {
           // Extract bits
-          auto wf = mxfp4_extract_bits_simd<S>(w_local);
+          auto wf = fp_extract_bits_simd<S, bits>(w_local);
           w_local += packs_per_simd;
           simd::Simd<float, S> x_simd = simd::load<T, S>(x_local);
           g_acc = g_acc + x_simd * wf;
@@ -585,8 +594,8 @@ void mxfp4_qmm_t_simd(
   }
 }
 
-template <typename T>
-void mxfp4_qmm_dispatch_transpose(
+template <typename T, int group_size, int bits>
+void fp_qmm_dispatch_transpose(
     T* result,
     const T* x,
     const uint32_t* w,
@@ -598,17 +607,17 @@ void mxfp4_qmm_dispatch_transpose(
   if (transposed_w) {
     // the simd size must be a multiple of the number of elements per word
     if constexpr (simd::max_size<T> % 8 == 0) {
-      mxfp4_qmm_t_simd<T>(result, x, w, scales, M, N, K);
+      fp_qmm_t_simd<T, group_size, bits>(result, x, w, scales, M, N, K);
     } else {
-      mxfp4_qmm_t<T>(result, x, w, scales, M, N, K);
+      fp_qmm_t<T, group_size, bits>(result, x, w, scales, M, N, K);
     }
   } else {
-    mxfp4_qmm<T>(result, x, w, scales, M, N, K);
+    fp_qmm<T, group_size, bits>(result, x, w, scales, M, N, K);
   }
 }
 
-template <typename T>
-void mxfp4_qmm_dispatch_typed(
+template <typename T, int group_size, int bits>
+void fp_qmm_dispatch_mode(
     array& out,
     const array& x,
     const array& w,
@@ -626,7 +635,7 @@ void mxfp4_qmm_dispatch_typed(
   auto w_ptr = w.data<uint32_t>();
   auto scales_ptr = scales.data<uint8_t>();
   for (int i = 0; i < batch_size; i++) {
-    mxfp4_qmm_dispatch_transpose<T>(
+    fp_qmm_dispatch_transpose<T, group_size, bits>(
         out_ptr + i * M * N,
         x_ptr + elem_to_loc(i * M * K, x.shape(), x.strides()),
         w_ptr + elem_to_loc(i * w_els, w.shape(), w.strides()),
@@ -638,21 +647,44 @@ void mxfp4_qmm_dispatch_typed(
   }
 }
 
-void mxfp4_qmm_dispatch(
+template <typename T>
+void fp_qmm_dispatch_typed(
     array& out,
     const array& x,
     const array& w,
     const array& scales,
+    int group_size,
+    int bits,
+    bool transposed_w) {
+  if (bits == 8) {
+    fp_qmm_dispatch_mode<T, 32, 8>(out, x, w, scales, transposed_w);
+  } else if (group_size == 32) {
+    fp_qmm_dispatch_mode<T, 32, 4>(out, x, w, scales, transposed_w);
+  } else {
+    fp_qmm_dispatch_mode<T, 16, 4>(out, x, w, scales, transposed_w);
+  }
+}
+
+void fp_qmm_dispatch(
+    array& out,
+    const array& x,
+    const array& w,
+    const array& scales,
+    int group_size,
+    int bits,
     bool transposed_w) {
   switch (x.dtype()) {
     case bfloat16:
-      mxfp4_qmm_dispatch_typed<bfloat16_t>(out, x, w, scales, transposed_w);
+      fp_qmm_dispatch_typed<bfloat16_t>(
+          out, x, w, scales, group_size, bits, transposed_w);
       break;
     case float16:
-      mxfp4_qmm_dispatch_typed<float16_t>(out, x, w, scales, transposed_w);
+      fp_qmm_dispatch_typed<float16_t>(
+          out, x, w, scales, group_size, bits, transposed_w);
       break;
     case float32:
-      mxfp4_qmm_dispatch_typed<float>(out, x, w, scales, transposed_w);
+      fp_qmm_dispatch_typed<float>(
+          out, x, w, scales, group_size, bits, transposed_w);
       break;
     default:
       throw std::invalid_argument(
@@ -765,9 +797,8 @@ void _bs_qmm_dispatch(
           "[quantized_matmul] only floating types are supported");
   }
 }
-
-template <typename T>
-void mxfp4_bs_qmm_dispatch_typed(
+template <typename T, int group_size, int bits>
+void fp_bs_qmm_dispatch_mode(
     array& out,
     const array& x,
     const array& w,
@@ -794,7 +825,7 @@ void mxfp4_bs_qmm_dispatch_typed(
         i, lhs_indices.shape(), lhs_indices.strides())];
     int w_idx = rhs_indices_ptr[elem_to_loc(
         i, rhs_indices.shape(), rhs_indices.strides())];
-    mxfp4_qmm_dispatch_transpose<T>(
+    fp_qmm_dispatch_transpose<T, group_size, bits>(
         out_ptr + i * M * N,
         x_ptr + elem_to_loc(x_idx * M * K, x.shape(), x.strides()),
         w_ptr + elem_to_loc(w_idx * w_els, w.shape(), w.strides()),
@@ -807,26 +838,75 @@ void mxfp4_bs_qmm_dispatch_typed(
   }
 }
 
-void mxfp4_bs_qmm_dispatch(
+template <typename T>
+void fp_bs_qmm_dispatch_typed(
     array& out,
     const array& x,
     const array& w,
     const array& scales,
     const array& lhs_indices,
     const array& rhs_indices,
+    int group_size,
+    int bits,
+    bool transposed_w) {
+  if (bits == 8) {
+    fp_bs_qmm_dispatch_mode<T, 32, 8>(
+        out, x, w, scales, lhs_indices, rhs_indices, transposed_w);
+  } else if (group_size == 32) {
+    fp_bs_qmm_dispatch_mode<T, 32, 4>(
+        out, x, w, scales, lhs_indices, rhs_indices, transposed_w);
+  } else {
+    fp_bs_qmm_dispatch_mode<T, 16, 4>(
+        out, x, w, scales, lhs_indices, rhs_indices, transposed_w);
+  }
+}
+
+void fp_bs_qmm_dispatch(
+    array& out,
+    const array& x,
+    const array& w,
+    const array& scales,
+    const array& lhs_indices,
+    const array& rhs_indices,
+    int group_size,
+    int bits,
     bool transposed_w) {
   switch (x.dtype()) {
     case float32:
-      mxfp4_bs_qmm_dispatch_typed<float>(
-          out, x, w, scales, lhs_indices, rhs_indices, transposed_w);
+      fp_bs_qmm_dispatch_typed<float>(
+          out,
+          x,
+          w,
+          scales,
+          lhs_indices,
+          rhs_indices,
+          group_size,
+          bits,
+          transposed_w);
       break;
     case float16:
-      mxfp4_bs_qmm_dispatch_typed<float16_t>(
-          out, x, w, scales, lhs_indices, rhs_indices, transposed_w);
+      fp_bs_qmm_dispatch_typed<float16_t>(
+          out,
+          x,
+          w,
+          scales,
+          lhs_indices,
+          rhs_indices,
+          group_size,
+          bits,
+          transposed_w);
       break;
     case bfloat16:
-      mxfp4_bs_qmm_dispatch_typed<bfloat16_t>(
-          out, x, w, scales, lhs_indices, rhs_indices, transposed_w);
+      fp_bs_qmm_dispatch_typed<bfloat16_t>(
+          out,
+          x,
+          w,
+          scales,
+          lhs_indices,
+          rhs_indices,
+          group_size,
+          bits,
+          transposed_w);
       break;
     default:
       throw std::invalid_argument(
@@ -881,8 +961,10 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
                       x = array::unsafe_weak_copy(x),
                       w = array::unsafe_weak_copy(w),
                       scales = array::unsafe_weak_copy(scales),
+                      group_size_ = group_size_,
+                      bits_ = bits_,
                       transpose_ = transpose_]() mutable {
-      mxfp4_qmm_dispatch(out, x, w, scales, transpose_);
+      fp_qmm_dispatch(out, x, w, scales, group_size_, bits_, transpose_);
     });
   }
 }
@@ -953,9 +1035,19 @@ void GatherQMM::eval_cpu(const std::vector<array>& inputs, array& out) {
                       scales = array::unsafe_weak_copy(scales),
                       lhs_indices = array::unsafe_weak_copy(lhs_indices),
                       rhs_indices = array::unsafe_weak_copy(rhs_indices),
+                      group_size_ = group_size_,
+                      bits_ = bits_,
                       transpose_ = transpose_]() mutable {
-      mxfp4_bs_qmm_dispatch(
-          out, x, w, scales, lhs_indices, rhs_indices, transpose_);
+      fp_bs_qmm_dispatch(
+          out,
+          x,
+          w,
+          scales,
+          lhs_indices,
+          rhs_indices,
+          group_size_,
+          bits_,
+          transpose_);
     });
   }
 }
@@ -1145,4 +1237,7 @@ void fast::ConvertFP8::eval_cpu(
   });
 }
 
+void QQMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
+  throw std::runtime_error("QQMatmul not implemented on CPU.");
+}
 } // namespace mlx::core

mlx/backend/cpu/simd/type.h

@@ -3,5 +3,9 @@
 #include "mlx/backend/cpu/simd/base_simd.h"
 
 #ifdef MLX_USE_ACCELERATE
+#if defined(__x86_64__)
+// the accelerate_simd implementation require neon -- use base implementation
+#else
 #include "mlx/backend/cpu/simd/accelerate_simd.h"
 #endif
+#endif

mlx/backend/cpu/svd.cpp

@@ -8,6 +8,183 @@
 
 namespace mlx::core {
 
+template <typename T, class Enable = void>
+struct SVDWork {};
+
+template <typename T>
+struct SVDWork<
+    T,
+    typename std::enable_if<std::is_floating_point<T>::value>::type> {
+  using R = T;
+
+  int N;
+  int M;
+  int K;
+  int lda;
+  int ldu;
+  int ldvt;
+  char jobz;
+  std::vector<array::Data> buffers;
+  int lwork;
+
+  SVDWork(int N, int M, int K, char jobz)
+      : N(N), M(M), K(K), lda(N), ldu(N), ldvt(M), jobz(jobz) {
+    T workspace_dimension = 0;
+
+    // Will contain the indices of eigenvectors that failed to converge (not
+    // used here but required by lapack).
+    buffers.emplace_back(allocator::malloc(sizeof(int) * 8 * K));
+
+    int lwork_query = -1;
+    int info;
+
+    // Compute workspace size.
+    gesdd<T>(
+        /* jobz = */ &jobz,
+        // M and N are swapped since lapack expects column-major.
+        /* m = */ &N,
+        /* n = */ &M,
+        /* a = */ nullptr,
+        /* lda = */ &lda,
+        /* s = */ nullptr,
+        /* u = */ nullptr,
+        /* ldu = */ &ldu,
+        /* vt = */ nullptr,
+        /* ldvt = */ &ldvt,
+        /* work = */ &workspace_dimension,
+        /* lwork = */ &lwork_query,
+        /* iwork = */ static_cast<int*>(buffers[0].buffer.raw_ptr()),
+        /* info = */ &info);
+
+    if (info != 0) {
+      std::stringstream ss;
+      ss << "[SVD::eval_cpu] workspace calculation failed with code " << info;
+      throw std::runtime_error(ss.str());
+    }
+
+    lwork = workspace_dimension;
+    buffers.emplace_back(allocator::malloc(sizeof(T) * lwork));
+  }
+
+  void run(T* a, R* s, T* u, T* vt) {
+    int info;
+    gesdd<T>(
+        /* jobz = */ &jobz,
+        // M and N are swapped since lapack expects column-major.
+        /* m = */ &N,
+        /* n = */ &M,
+        /* a = */ a,
+        /* lda = */ &lda,
+        /* s = */ s,
+        // According to the identity above, lapack will write Vᵀᵀ as U.
+        /* u = */ u,
+        /* ldu = */ &ldu,
+        // According to the identity above, lapack will write Uᵀ as Vᵀ.
+        /* vt = */ vt,
+        /* ldvt = */ &ldvt,
+        /* work = */ static_cast<T*>(buffers[1].buffer.raw_ptr()),
+        /* lwork = */ &lwork,
+        /* iwork = */ static_cast<int*>(buffers[0].buffer.raw_ptr()),
+        /* info = */ &info);
+
+    if (info != 0) {
+      std::stringstream ss;
+      ss << "svd_impl: sgesvdx_ failed with code " << info;
+      throw std::runtime_error(ss.str());
+    }
+  }
+};
+
+template <>
+struct SVDWork<std::complex<float>> {
+  using T = std::complex<float>;
+  using R = float;
+
+  int N;
+  int M;
+  int K;
+  int lda;
+  int ldu;
+  int ldvt;
+  char jobz;
+  std::vector<array::Data> buffers;
+  int lwork;
+
+  SVDWork(int N, int M, int K, char jobz)
+      : N(N), M(M), K(K), lda(N), ldu(N), ldvt(M), jobz(jobz) {
+    T workspace_dimension = 0;
+
+    // Will contain the indices of eigenvectors that failed to converge (not
+    // used here but required by lapack).
+    buffers.emplace_back(allocator::malloc(sizeof(int) * 8 * K));
+
+    const int lrwork =
+        jobz == 'A' ? std::max(1, 5 * K * K + 5 * K) : std::max(1, 7 * K);
+    buffers.emplace_back(allocator::malloc(sizeof(float) * lrwork));
+
+    int lwork_query = -1;
+    int work_query = -1;
+    int info;
+
+    // Compute workspace size.
+    gesdd<T>(
+        /* jobz = */ &jobz,
+        // M and N are swapped since lapack expects column-major.
+        /* m = */ &N,
+        /* n = */ &M,
+        /* a = */ nullptr,
+        /* lda = */ &lda,
+        /* s = */ nullptr,
+        /* u = */ nullptr,
+        /* ldu = */ &ldu,
+        /* vt = */ nullptr,
+        /* ldvt = */ &ldvt,
+        /* work = */ &workspace_dimension,
+        /* lwork = */ &lwork_query,
+        /* rwork = */ static_cast<float*>(buffers[1].buffer.raw_ptr()),
+        /* iwork = */ static_cast<int*>(buffers[0].buffer.raw_ptr()),
+        /* info = */ &info);
+
+    if (info != 0) {
+      std::stringstream ss;
+      ss << "[SVD::eval_cpu] workspace calculation failed with code " << info;
+      throw std::runtime_error(ss.str());
+    }
+
+    lwork = workspace_dimension.real();
+    buffers.emplace_back(allocator::malloc(sizeof(T) * lwork));
+  }
+
+  void run(T* a, R* s, T* u, T* vt) {
+    int info;
+    gesdd<T>(
+        /* jobz = */ &jobz,
+        // M and N are swapped since lapack expects column-major.
+        /* m = */ &N,
+        /* n = */ &M,
+        /* a = */ a,
+        /* lda = */ &lda,
+        /* s = */ s,
+        // According to the identity above, lapack will write Vᵀᵀ as U.
+        /* u = */ u,
+        /* ldu = */ &ldu,
+        // According to the identity above, lapack will write Uᵀ as Vᵀ.
+        /* vt = */ vt,
+        /* ldvt = */ &ldvt,
+        /* work = */ static_cast<T*>(buffers[2].buffer.raw_ptr()),
+        /* lwork = */ &lwork,
+        /* rwork = */ static_cast<float*>(buffers[1].buffer.raw_ptr()),
+        /* iwork = */ static_cast<int*>(buffers[0].buffer.raw_ptr()),
+        /* info = */ &info);
+
+    if (info != 0) {
+      std::stringstream ss;
+      ss << "svd_impl: sgesvdx_ failed with code " << info;
+      throw std::runtime_error(ss.str());
+    }
+  }
+};
+
 template <typename T>
 void svd_impl(
     const array& a,
@@ -27,6 +204,8 @@ void svd_impl(
   const int N = a.shape(-1);
   const int K = std::min(M, N);
 
+  using R = typename SVDWork<T>::R;
+
   size_t num_matrices = a.size() / (M * N);
 
   // lapack clobbers the input, so we have to make a copy.
@@ -42,7 +221,7 @@ void svd_impl(
   encoder.set_input_array(a);
   auto in_ptr = in.data<T>();
   T* u_ptr;
-  T* s_ptr;
+  R* s_ptr;
   T* vt_ptr;
 
   if (compute_uv) {
@@ -58,7 +237,7 @@ void svd_impl(
     encoder.set_output_array(s);
     encoder.set_output_array(vt);
 
-    s_ptr = s.data<T>();
+    s_ptr = s.data<R>();
     u_ptr = u.data<T>();
     vt_ptr = vt.data<T>();
   } else {
@@ -68,96 +247,26 @@ void svd_impl(
 
     encoder.set_output_array(s);
 
-    s_ptr = s.data<T>();
+    s_ptr = s.data<R>();
     u_ptr = nullptr;
     vt_ptr = nullptr;
   }
 
   encoder.dispatch([in_ptr, u_ptr, s_ptr, vt_ptr, M, N, K, num_matrices]() {
-    // A of shape M x N. The leading dimension is N since lapack receives Aᵀ.
-    const int lda = N;
-    // U of shape M x M. (N x N in lapack).
-    const int ldu = N;
-    // Vᵀ of shape N x N. (M x M in lapack).
-    const int ldvt = M;
-
-    auto jobz = (u_ptr) ? "A" : "N";
-
-    T workspace_dimension = 0;
-
-    // Will contain the indices of eigenvectors that failed to converge (not
-    // used here but required by lapack).
-    auto iwork = array::Data{allocator::malloc(sizeof(int) * 8 * K)};
-
-    static const int lwork_query = -1;
-
-    int info;
-
-    // Compute workspace size.
-    gesdd<T>(
-        /* jobz = */ jobz,
-        // M and N are swapped since lapack expects column-major.
-        /* m = */ &N,
-        /* n = */ &M,
-        /* a = */ nullptr,
-        /* lda = */ &lda,
-        /* s = */ nullptr,
-        /* u = */ nullptr,
-        /* ldu = */ &ldu,
-        /* vt = */ nullptr,
-        /* ldvt = */ &ldvt,
-        /* work = */ &workspace_dimension,
-        /* lwork = */ &lwork_query,
-        /* iwork = */ static_cast<int*>(iwork.buffer.raw_ptr()),
-        /* info = */ &info);
-
-    if (info != 0) {
-      std::stringstream ss;
-      ss << "[SVD::eval_cpu] workspace calculation failed with code " << info;
-      throw std::runtime_error(ss.str());
-    }
-
-    const int lwork = workspace_dimension;
-    auto scratch = array::Data{allocator::malloc(sizeof(T) * lwork)};
-
+    auto jobz = (u_ptr) ? 'A' : 'N';
+    SVDWork<T> svd_work(N, M, K, jobz);
     // Loop over matrices.
     for (int i = 0; i < num_matrices; i++) {
-      gesdd<T>(
-          /* jobz = */ jobz,
-          // M and N are swapped since lapack expects column-major.
-          /* m = */ &N,
-          /* n = */ &M,
-          /* a = */ in_ptr + M * N * i,
-          /* lda = */ &lda,
-          /* s = */ s_ptr + K * i,
-          // According to the identity above, lapack will write Vᵀᵀ as U.
-          /* u = */ vt_ptr ? vt_ptr + N * N * i : nullptr,
-          /* ldu = */ &ldu,
-          // According to the identity above, lapack will write Uᵀ as Vᵀ.
-          /* vt = */ u_ptr ? u_ptr + M * M * i : nullptr,
-          /* ldvt = */ &ldvt,
-          /* work = */ static_cast<T*>(scratch.buffer.raw_ptr()),
-          /* lwork = */ &lwork,
-          /* iwork = */ static_cast<int*>(iwork.buffer.raw_ptr()),
-          /* info = */ &info);
-
-      if (info != 0) {
-        std::stringstream ss;
-        ss << "svd_impl: sgesvdx_ failed with code " << info;
-        throw std::runtime_error(ss.str());
-      }
+      svd_work.run(
+          in_ptr + M * N * i,
+          s_ptr + K * i,
+          vt_ptr ? vt_ptr + N * N * i : nullptr,
+          u_ptr ? u_ptr + M * M * i : nullptr);
     }
   });
   encoder.add_temporary(in);
 }
 
-template <typename T>
-void compute_svd(
-    const array& a,
-    bool compute_uv,
-    std::vector<array>& outputs,
-    Stream stream) {}
-
 void SVD::eval_cpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
@@ -168,9 +277,12 @@ void SVD::eval_cpu(
     case float64:
       svd_impl<double>(inputs[0], outputs, compute_uv_, stream());
       break;
+    case complex64:
+      svd_impl<std::complex<float>>(inputs[0], outputs, compute_uv_, stream());
+      break;
     default:
       throw std::runtime_error(
-          "[SVD::eval_cpu] only supports float32 or float64.");
+          "[SVD::eval_cpu] only supports float32, float64, or complex64.");
   }
 }
 

mlx/backend/cuda/CMakeLists.txt

@@ -18,6 +18,7 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_conv.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_grouped_conv.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/cublas_utils.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/cuda.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/cudnn_utils.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/custom_kernel.cpp
@@ -28,6 +29,7 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/fence.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/gemms/gemv.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/gemms/cublas_gemm.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/gemms/grouped_gemm_unaligned.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/jit_module.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/kernel_utils.cu
@@ -64,6 +66,12 @@ add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/unary)
 # fp4 is not available on < 12.8
 if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 12.8.0)
   target_include_directories(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/)
+else()
+  target_sources(
+    mlx
+    PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm.cpp
+            ${CMAKE_CURRENT_SOURCE_DIR}/quantized/cublas_qqmm.cpp
+            ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_utils.cu)
 endif()
 
 if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.9.0)
@@ -74,8 +82,6 @@ else()
     mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/gemms/cublas_gemm_batched_12_0.cpp)
 endif()
 
-target_compile_definitions(mlx PRIVATE MLX_USE_CUDA)
-
 # Embed kernel sources in binary for JIT compilation.
 file(
   GLOB MLX_JIT_SOURCES
@@ -94,6 +100,10 @@ add_custom_target(cuda_jit_sources DEPENDS gen/cuda_jit_sources.h)
 add_dependencies(mlx cuda_jit_sources)
 target_include_directories(mlx PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/gen")
 
+# ------------------------ Compilation configs ------------------------
+
+target_compile_definitions(mlx PRIVATE MLX_USE_CUDA)
+
 # Enable defining device lambda functions.
 target_compile_options(mlx
                        PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--extended-lambda>")
@@ -116,6 +126,10 @@ endif()
 target_compile_options(
   mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--Wno-deprecated-gpu-targets>")
 
+# Suppress nvcc warnings on MLX headers.
+target_compile_options(mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe
+                                   --diag_suppress=997>)
+
 # Use stronger binaries compression. This feature was introduced in CUDA 12.8
 # and requires drivers released after CUDA 12.4.
 if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8.0)
@@ -123,26 +137,59 @@ if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8.0)
     mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--compress-mode=size>")
 endif()
 
-# Compute capability >= 7.0 is required for synchronization between CPU/GPU with
-# managed memory.
+# Use native CUDA arch by default.
 if(NOT DEFINED MLX_CUDA_ARCHITECTURES)
   execute_process(
-    COMMAND bash detect_cuda_arch.sh
-    WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+    COMMAND __nvcc_device_query
     OUTPUT_VARIABLE MLX_CUDA_ARCHITECTURES
     OUTPUT_STRIP_TRAILING_WHITESPACE)
+  set(UPGRADABLE_ARCHITECTURES "90;100;121")
+  if(MLX_CUDA_ARCHITECTURES STREQUAL "")
+    message(
+      FATAL_ERROR
+        "Can not get native CUDA arch, must set MLX_CUDA_ARCHITECTURES")
+  elseif(MLX_CUDA_ARCHITECTURES IN_LIST UPGRADABLE_ARCHITECTURES)
+    # Use arch-specific compute capability whenever possible.
+    set(MLX_CUDA_ARCHITECTURES "${MLX_CUDA_ARCHITECTURES}a")
+  endif()
 endif()
 message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}")
 set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES
                                      "${MLX_CUDA_ARCHITECTURES}")
 
+if(MLX_BUILD_PYTHON_BINDINGS)
+  set_property(
+    TARGET mlx
+    APPEND
+    PROPERTY INSTALL_RPATH
+             # The paths here should match the install_requires in setup.py.
+             "$ORIGIN/../../nvidia/cublas/lib"
+             "$ORIGIN/../../nvidia/cuda_nvrtc/lib"
+             "$ORIGIN/../../nvidia/cudnn/lib"
+             "$ORIGIN/../../nvidia/nccl/lib")
+endif()
+
+# ------------------------ Dependencies ------------------------
+
 # Use fixed version of CCCL.
 FetchContent_Declare(
   cccl
   URL "https://github.com/NVIDIA/cccl/releases/download/v2.8.1/cccl-v2.8.1.zip")
 FetchContent_MakeAvailable(cccl)
 target_include_directories(mlx BEFORE PRIVATE "${cccl_SOURCE_DIR}/include")
-set_target_properties(mlx PROPERTIES CCCL_DIR "${cccl_SOURCE_DIR}/include")
+
+# Install CCCL headers for JIT.
+install(DIRECTORY ${cccl_SOURCE_DIR}/include/cuda
+        DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/cccl)
+install(DIRECTORY ${cccl_SOURCE_DIR}/include/nv
+        DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/cccl)
+
+# The binary of C++ tests will not be installed so it can not find the CCCL
+# headers, and we have to hard-code the path.
+if(MLX_BUILD_TESTS)
+  target_compile_definitions(mlx
+                             PRIVATE MLX_CCCL_DIR="${cccl_SOURCE_DIR}/include")
+endif()
 
 # Use fixed version of NVTX.
 FetchContent_Declare(
@@ -181,9 +228,13 @@ target_link_libraries(mlx PRIVATE cudnn_frontend)
 include(${cudnn_frontend_SOURCE_DIR}/cmake/cuDNN.cmake)
 target_link_libraries(mlx PRIVATE CUDNN::cudnn_all)
 
-# Suppress nvcc warnings on MLX headers.
-target_compile_options(mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe
-                                   --diag_suppress=997>)
-# Install CCCL headers for JIT.
-install(DIRECTORY ${cccl_SOURCE_DIR}/include/cuda
-        DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/cccl)
+# Use header-only CUTLASS.
+FetchContent_Declare(
+  cutlass
+  GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
+  GIT_TAG v4.3.2
+  GIT_SHALLOW TRUE
+  SOURCE_SUBDIR include EXCLUDE_FROM_ALL)
+FetchContent_MakeAvailable(cutlass)
+target_include_directories(
+  mlx PRIVATE $<BUILD_INTERFACE:${cutlass_SOURCE_DIR}/include>)

mlx/backend/cuda/allocator.cpp

@@ -20,6 +20,19 @@ constexpr int page_size = 16384;
 // Any allocations smaller than this will try to use the small pool
 constexpr int small_block_size = 8;
 
+#if CUDART_VERSION >= 13000
+inline cudaMemLocation cuda_mem_loc(int i) {
+  cudaMemLocation loc;
+  loc.type = cudaMemLocationTypeDevice;
+  loc.id = i;
+  return loc;
+}
+#else
+inline int cuda_mem_loc(int i) {
+  return i;
+}
+#endif // CUDART_VERSION >= 13000
+
 // The small pool size in bytes. This should be a multiple of the host page
 // size and small_block_size.
 constexpr int small_pool_size = 4 * page_size;
@@ -35,13 +48,7 @@ SmallSizePool::SmallSizePool() {
   int device_count = 0;
   CHECK_CUDA_ERROR(cudaGetDeviceCount(&device_count));
   for (int i = 0; i < device_count; ++i) {
-#if CUDART_VERSION >= 13000
-    cudaMemLocation loc;
-    loc.type = cudaMemLocationTypeDevice;
-    loc.id = i;
-#else
-    int loc = i;
-#endif // CUDART_VERSION >= 13000
+    auto loc = cuda_mem_loc(i);
     CHECK_CUDA_ERROR(
         cudaMemAdvise(data_, small_pool_size, cudaMemAdviseSetAccessedBy, loc));
   }
@@ -90,9 +97,10 @@ CudaAllocator::CudaAllocator()
           page_size,
           [](CudaBuffer* buf) { return buf->size; },
           [this](CudaBuffer* buf) { cuda_free(buf); }) {
-  size_t free, total;
-  CHECK_CUDA_ERROR(cudaMemGetInfo(&free, &total));
-  memory_limit_ = total * 0.95;
+  size_t free;
+  CHECK_CUDA_ERROR(cudaMemGetInfo(&free, &total_memory_));
+  memory_limit_ = total_memory_ * 0.95;
+  free_limit_ = total_memory_ - memory_limit_;
   max_pool_size_ = memory_limit_;
 
   int device_count = 0;
@@ -104,6 +112,10 @@ CudaAllocator::CudaAllocator()
     cudaStream_t s;
     CHECK_CUDA_ERROR(cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking));
     free_streams_.push_back(s);
+
+    cudaMemPool_t mem_pool;
+    CHECK_CUDA_ERROR(cudaDeviceGetDefaultMemPool(&mem_pool, i));
+    mem_pools_.push_back(mem_pool);
   }
   CHECK_CUDA_ERROR(cudaSetDevice(curr));
 }
@@ -154,19 +166,35 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
     }
     lock.unlock();
     if (!buf) {
-      buf = new CudaBuffer{nullptr, size, device};
-      cudaError_t err;
+      void* data = nullptr;
       if (device == -1) {
-        err = cudaMallocManaged(&buf->data, size);
+        CHECK_CUDA_ERROR(cudaMallocManaged(&data, size));
       } else {
-        err = cudaMallocAsync(&buf->data, size, stream);
+        CHECK_CUDA_ERROR(cudaMallocAsync(&data, size, stream));
       }
-      if (err != cudaSuccess && err != cudaErrorMemoryAllocation) {
-        throw std::runtime_error(fmt::format(
-            "cudaMallocManaged failed: {}.", cudaGetErrorString(err)));
+      if (!data) {
+        std::ostringstream msg;
+        msg << "[malloc] Unable to allocate " << size << " bytes.";
+        throw std::runtime_error(msg.str());
       }
+      buf = new CudaBuffer{data, size, device};
     }
     lock.lock();
+
+    // If any cuda memory pool has too much reserved memory, clear some
+    // memory from the cache. This prevents graph / kernel execution failing
+    // from OOM
+    if (get_cache_memory() > 0) {
+      for (auto p : mem_pools_) {
+        size_t used = 0;
+        CHECK_CUDA_ERROR(cudaMemPoolGetAttribute(
+            p, cudaMemPoolAttrReservedMemCurrent, &used));
+        if (used > (total_memory_ - free_limit_)) {
+          buffer_cache_.release_cached_buffers(free_limit_);
+          break;
+        }
+      }
+    }
   }
   active_memory_ += buf->size;
   peak_memory_ = std::max(active_memory_, peak_memory_);
@@ -176,7 +204,7 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
     buffer_cache_.release_cached_buffers(get_cache_memory() - max_pool_size_);
   }
   // Copy to managed here if the buffer is not on the right device
-  if (buf->device != device) {
+  if (buf->device >= 0 && buf->device != device) {
     copy_to_managed(*buf);
   }
   return Buffer{buf};
@@ -219,9 +247,9 @@ void CudaAllocator::cuda_free(CudaBuffer* buf) {
     scalar_pool_.free(buf);
   } else {
     if (buf->device >= 0) {
-      cudaFreeAsync(buf->data, free_streams_[buf->device]);
+      CHECK_CUDA_ERROR(cudaFreeAsync(buf->data, free_streams_[buf->device]));
     } else {
-      cudaFree(buf->data);
+      CHECK_CUDA_ERROR(cudaFree(buf->data));
     }
     delete buf;
   }

mlx/backend/cuda/allocator.h

@@ -71,11 +71,14 @@ class CudaAllocator : public allocator::Allocator {
 
   std::mutex mutex_;
   size_t memory_limit_;
+  size_t free_limit_;
+  size_t total_memory_;
   size_t max_pool_size_;
   BufferCache<CudaBuffer> buffer_cache_;
   size_t active_memory_{0};
   size_t peak_memory_{0};
   std::vector<cudaStream_t> free_streams_;
+  std::vector<cudaMemPool_t> mem_pools_;
   SmallSizePool scalar_pool_;
 };
 

mlx/backend/cuda/conv.cpp

@@ -15,19 +15,16 @@ namespace mlx::core {
 
 namespace {
 
-// Alias for better readability.
-#define CONV_FORWARD CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR
-#define CONV_BACKWARD_INPUT \
-  CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR
-#define CONV_BACKWARD_WEIGHT \
-  CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR
-
-// Custom placeholder representing fallback kernel.
-#define CONV_FALLBACK static_cast<cudnnBackendDescriptorType_t>(-1)
+enum ConvBackendType {
+  CONV_FALLBACK,
+  CONV_FORWARD,
+  CONV_BACKWARD_INPUT,
+  CONV_BACKWARD_WEIGHT,
+};
 
 struct ConvCacheKey {
   int device_id;
-  cudnnDataType_t cudnn_dtype;
+  fe::DataType_t cudnn_dtype;
   std::array<int, MAX_NDIM> input_shape;
   std::array<int, MAX_NDIM> weight_shape;
   std::array<int, MAX_NDIM> stride;
@@ -44,15 +41,13 @@ struct ConvCacheKey {
 auto& conv_cache() {
   static LRUBytesKeyCache<
       ConvCacheKey,
-      std::pair<
-          cudnnBackendDescriptorType_t,
-          std::optional<cudnn_frontend::ExecutionPlan>>>
+      std::pair<ConvBackendType, std::optional<DnnGraph>>>
       cache("MLX_CUDA_CONV_CACHE_SIZE", /* default_capacity */ 128);
   return cache;
 }
 
-auto get_conv_op_settings(
-    cudnnBackendDescriptorType_t backend_type,
+auto get_conv_settings(
+    ConvBackendType backend_type,
     array& x,
     array& w,
     array& y,
@@ -68,8 +63,8 @@ auto get_conv_op_settings(
     for (int i = 0; i < padding_lo.size(); ++i) {
       int wt_size = 1 + kernel_dilation[i] * (w.shape(1 + i) - 1);
       padding_lo[i] = wt_size - padding_lo[i] - 1;
-      int in_size = 1 + kernel_strides[i] * (x.shape(1 + i) - 1);
-      int out_size = 1 + input_dilation[i] * (y.shape(1 + i) - 1);
+      int in_size = 1 + kernel_strides[i] * (y.shape(1 + i) - 1);
+      int out_size = 1 + input_dilation[i] * (x.shape(1 + i) - 1);
       padding_hi[i] = out_size - in_size + padding_hi[i];
     }
     return std::make_tuple(
@@ -95,49 +90,57 @@ auto get_conv_op_settings(
   }
 }
 
-std::optional<cudnn_frontend::OperationGraph> build_conv_op_graph(
+std::optional<DnnGraph> build_conv_graph(
     cu::CommandEncoder& encoder,
-    cudnnBackendDescriptorType_t backend_type,
+    ConvBackendType backend_type,
     Dtype dtype,
     array& x,
     array& w,
     array& y,
-    const SmallVector<int64_t>& stride,
-    const SmallVector<int64_t>& padding_lo,
-    const SmallVector<int64_t>& padding_hi,
-    const SmallVector<int64_t>& dilation) {
-  try {
-    auto compute_dtype = (dtype == float16 || dtype == bfloat16)
-        ? CUDNN_DATA_FLOAT
-        : dtype_to_cudnn_type(dtype);
-    auto conv_desc = cudnn_frontend::ConvDescBuilder()
-                         .setDataType(compute_dtype)
-                         .setMathMode(CUDNN_CROSS_CORRELATION)
-                         .setNDims(stride.size())
-                         .setStrides(stride.size(), stride.data())
-                         .setPrePadding(padding_lo.size(), padding_lo.data())
-                         .setPostPadding(padding_hi.size(), padding_hi.data())
-                         .setDilation(dilation.size(), dilation.data())
-                         .build();
+    const std::vector<int64_t>& stride,
+    const std::vector<int64_t>& padding_lo,
+    const std::vector<int64_t>& padding_hi,
+    const std::vector<int64_t>& dilation) {
+  auto compute_dtype =
+      (dtype == float16 || dtype == bfloat16) ? float32 : dtype;
+  DnnGraph graph(encoder.device().cudnn_handle(), dtype, compute_dtype);
+  auto x_ = graph.tensor_nchw("X", 'x', x);
+  auto w_ = graph.tensor_nchw("W", 'w', w);
 
-    auto op = cudnn_frontend::OperationBuilder(backend_type)
-                  .setxDesc(build_cudnn_tensor_nchw('x', x))
-                  .setwDesc(build_cudnn_tensor_nchw('w', w))
-                  .setyDesc(build_cudnn_tensor_nchw('y', y))
-                  .setcDesc(conv_desc)
-                  .build();
+  auto set_options = [&](auto& options) {
+    options.set_compute_data_type(dtype_to_cudnn_type(compute_dtype))
+        .set_convolution_mode(fe::ConvolutionMode_t::CROSS_CORRELATION)
+        .set_stride(stride)
+        .set_pre_padding(padding_lo)
+        .set_post_padding(padding_hi)
+        .set_dilation(dilation);
+  };
 
-    std::array<cudnn_frontend::Operation const*, 1> ops = {&op};
-    return cudnn_frontend::OperationGraphBuilder()
-        .setHandle(encoder.device().cudnn_handle())
-        .setOperationGraph(ops.size(), ops.data())
-        .build();
-  } catch (cudnn_frontend::cudnnException& error) {
-    if (error.getCudnnStatus() != CUDNN_STATUS_BAD_PARAM) {
-      throw;
-    }
+  std::shared_ptr<fe::graph::Tensor_attributes> y_;
+  if (backend_type == CONV_FORWARD) {
+    auto options = fe::graph::Conv_fprop_attributes();
+    set_options(options);
+    y_ = graph.conv_fprop(x_, w_, options);
+  } else if (backend_type == CONV_BACKWARD_INPUT) {
+    auto options = fe::graph::Conv_dgrad_attributes();
+    set_options(options);
+    y_ = graph.conv_dgrad(x_, w_, options);
+  } else if (backend_type == CONV_BACKWARD_WEIGHT) {
+    auto options = fe::graph::Conv_wgrad_attributes();
+    set_options(options);
+    y_ = graph.conv_wgrad(w_, x_, options);
+  }
+  graph.tensor_nchw(y_, 'y', y)->set_output(true);
+
+  if (graph.prepare().is_bad()) {
     return std::nullopt;
   }
+  graph.deselect_numeric_notes({fe::NumericalNote_t::DOWN_CONVERT_INPUTS});
+  if (dtype == float32 && !env::enable_tf32()) {
+    graph.deselect_numeric_notes({fe::NumericalNote_t::TENSOR_CORE});
+  }
+  CHECK_CUDNN_FE_ERROR(graph.build());
+  return graph;
 }
 
 // Transpose from (C_out, H, W, C_in / groups) to (C_in, H, W, C_out / groups).
@@ -181,7 +184,7 @@ array group_transpose(
 // eval_gpu, with cost of possible redundant copies.
 std::tuple<array, array, array> prepare_args(
     cu::CommandEncoder& encoder,
-    cudnnBackendDescriptorType_t backend_type,
+    ConvBackendType backend_type,
     array in,
     array wt,
     array out,
@@ -221,27 +224,11 @@ std::tuple<array, array, array> prepare_args(
   return {std::move(in), std::move(wt), std::move(out)};
 }
 
-// Get the x/w/y args from the in/wt/out args depending on backend type.
-inline std::tuple<array&, array&, array&> dispatch_args(
-    cudnnBackendDescriptorType_t backend_type,
-    array& in,
-    array& wt,
-    array& out) {
-  switch (backend_type) {
-    case CONV_BACKWARD_INPUT:
-      return {out, wt, in};
-    case CONV_BACKWARD_WEIGHT:
-      return {in, out, wt};
-    default:
-      return {in, wt, out};
-  }
-}
-
 // Register inputs and outputs before actually running conv op. Can only be
 // called once per eval_gpu.
 void register_args(
     cu::CommandEncoder& encoder,
-    cudnnBackendDescriptorType_t backend_type,
+    ConvBackendType backend_type,
     array& in,
     array& wt,
     array& intermediate_out,
@@ -297,16 +284,19 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
       get_alignment(wt),
       get_alignment(out)};
   if (auto it = conv_cache().find(cache_key); it != conv_cache().end()) {
-    auto& [backend_type, plan] = it->second;
-    if (plan) {
-      // Run cached plan.
+    auto& [backend_type, graph] = it->second;
+    if (graph) {
+      // Run cached graph.
       std::tie(in, wt, out) =
           prepare_args(encoder, backend_type, in, wt, out, groups_, s);
       register_args(encoder, backend_type, in, wt, out, out_);
-      auto [x, w, y] = dispatch_args(backend_type, in, wt, out);
-      if (!encode_cudnn_plan(encoder, *plan, {'x', 'w', 'y'}, x, w, y)) {
-        throw std::runtime_error("[conv] Cached plan failed to execute.");
-      }
+      CHECK_CUDNN_FE_ERROR(graph->encode_capturing(
+          encoder,
+          {
+              {'x', gpu_ptr<void>(in)},
+              {'w', gpu_ptr<void>(wt)},
+              {'y', gpu_ptr<void>(out)},
+          }));
     } else {
       // Run fallback kernel.
       gemm_conv(
@@ -327,7 +317,7 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
 
   // There is no reliable way to deduce the proper cuDNN backend for the
   // convolution, so we make a best guess and then try.
-  SmallVector<cudnnBackendDescriptorType_t, 2> try_backends;
+  SmallVector<ConvBackendType, 2> try_backends;
   if (flip_) {
     // When weight is flipped, we assume it is backward input convolution.
     try_backends.push_back(CONV_BACKWARD_INPUT);
@@ -345,13 +335,12 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
   }
 
   // Try to build op graph.
-  cudnnBackendDescriptorType_t backend_type;
-  std::optional<cudnn_frontend::OperationGraph> op_graph;
+  ConvBackendType backend_type;
+  std::optional<DnnGraph> graph;
   for (auto try_backend : try_backends) {
-    auto [in_copy, wt_copy, out_copy] =
+    auto [x, w, y] =
         prepare_args(encoder, try_backend, in, wt, out, groups_, s);
-    auto [x, w, y] = dispatch_args(try_backend, in_copy, wt_copy, out_copy);
-    auto [stride, padding_lo, padding_hi, dilation] = get_conv_op_settings(
+    auto [stride, padding_lo, padding_hi, dilation] = get_conv_settings(
         try_backend,
         x,
         w,
@@ -361,7 +350,7 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
         padding_hi_,
         kernel_dilation_,
         input_dilation_);
-    op_graph = build_conv_op_graph(
+    graph = build_conv_graph(
         encoder,
         try_backend,
         dtype,
@@ -372,30 +361,27 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
         padding_lo,
         padding_hi,
         dilation);
-    if (op_graph) {
+    if (graph) {
       backend_type = try_backend;
-      in = std::move(in_copy);
-      wt = std::move(wt_copy);
-      out = std::move(out_copy);
+      in = std::move(x);
+      wt = std::move(w);
+      out = std::move(y);
       break;
     }
   }
 
-  if (op_graph) {
-    // Find a plan for the graph and execute it.
-    auto plan = find_cudnn_plan_from_op_graph(
-        encoder.device().cudnn_handle(), backend_type, dtype, *op_graph);
-    if (plan) {
-      // Setup inputs and outputs.
-      register_args(encoder, backend_type, in, wt, out, out_);
-
-      auto [x, w, y] = dispatch_args(backend_type, in, wt, out);
-      if (encode_cudnn_plan(encoder, *plan, {'x', 'w', 'y'}, x, w, y)) {
-        conv_cache().emplace(
-            cache_key, std::make_pair(backend_type, std::move(*plan)));
-        return;
-      }
-    }
+  if (graph) {
+    register_args(encoder, backend_type, in, wt, out, out_);
+    CHECK_CUDNN_FE_ERROR(graph->encode_capturing(
+        encoder,
+        {
+            {'x', gpu_ptr<void>(in)},
+            {'w', gpu_ptr<void>(wt)},
+            {'y', gpu_ptr<void>(out)},
+        }));
+    conv_cache().emplace(
+        cache_key, std::make_pair(backend_type, std::move(*graph)));
+    return;
   }
 
   // Use fallback kernel for settings not supported by cuDNN.

mlx/backend/cuda/copy/copy_general_input.cu

@@ -5,6 +5,7 @@
 #include <cooperative_groups.h>
 
 namespace mlx::core {
+static constexpr int TILE_SIZE = 16;
 
 namespace cu {
 
@@ -73,6 +74,53 @@ __global__ void copy_g(
   store_vector(out + shape_x * index_rest, index_x, out_vec, shape_x);
 }
 
+template <typename In, typename Out, int N_READS>
+__global__ void
+copy_col_row(const In* in, Out* out, int64_t rows, int64_t cols) {
+  __shared__ Out
+      tile[N_READS * TILE_SIZE][N_READS * TILE_SIZE + 4 / sizeof(Out)];
+
+  auto block = cg::this_thread_block();
+  auto grid = cg::this_grid();
+
+  auto tile_row = grid.block_index().x * TILE_SIZE * N_READS;
+  auto tile_col = grid.block_index().y * TILE_SIZE * N_READS;
+
+  auto tidx = block.thread_index().x;
+  auto tidy = N_READS * block.thread_index().y;
+
+  auto in_ptr = in + (tile_col + tidy) * rows + tile_row;
+
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    if ((tile_col + tidy + i) < cols) {
+      auto in_vec = load_vector<N_READS>(in_ptr, tidx, rows - tile_row, In(0));
+#pragma unroll
+      for (int j = 0; j < N_READS; ++j) {
+        tile[N_READS * tidx + j][tidy + i] = CastOp<In, Out>{}(in_vec[j]);
+      }
+      in_ptr += rows;
+    }
+  }
+
+  block.sync();
+
+  auto out_ptr = out + (tile_row + tidy) * cols + tile_col;
+
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    if ((tile_row + tidy + i) < rows) {
+      AlignedVector<Out, N_READS> out_vec;
+#pragma unroll
+      for (int j = 0; j < N_READS; ++j) {
+        out_vec[j] = tile[tidy + i][N_READS * tidx + j];
+      }
+      store_vector(out_ptr, tidx, out_vec, cols - tile_col);
+      out_ptr += cols;
+    }
+  }
+}
+
 } // namespace cu
 
 void copy_general_input(
@@ -86,20 +134,46 @@ void copy_general_input(
     const Strides& strides_in) {
   dispatch_all_types(in.dtype(), [&](auto in_type_tag) {
     dispatch_all_types(out.dtype(), [&](auto out_type_tag) {
+      using InType = cuda_type_t<MLX_GET_TYPE(in_type_tag)>;
+      using OutType = cuda_type_t<MLX_GET_TYPE(out_type_tag)>;
+      const InType* in_ptr = gpu_ptr<InType>(in) + offset_in;
+      OutType* out_ptr = gpu_ptr<OutType>(out) + offset_out;
+      int ndim = shape.size();
+
+      // Column contiguous to row contiguous specialization
+      if (ndim == 2 && strides_in[0] == 1 && strides_in[1] == shape[0]) {
+        constexpr int work_per_thread =
+            std::min(static_cast<int>(16 / sizeof(OutType)), 8);
+        dim3 block_dims = {TILE_SIZE, TILE_SIZE};
+        uint32_t num_blocks_x =
+            cuda::ceil_div(shape[0], TILE_SIZE * work_per_thread);
+        uint32_t num_blocks_y =
+            cuda::ceil_div(shape[1], TILE_SIZE * work_per_thread);
+        auto kernel = cu::copy_col_row<InType, OutType, work_per_thread>;
+        encoder.add_kernel_node(
+            kernel,
+            {num_blocks_x, num_blocks_y},
+            block_dims,
+            0,
+            in_ptr,
+            out_ptr,
+            int64_t(shape[0]),
+            int64_t(shape[1]));
+        return;
+      }
+
       dispatch_bool(
           in.data_size() > INT32_MAX || out.data_size() > INT32_MAX,
           [&](auto large) {
-            using InType = cuda_type_t<MLX_GET_TYPE(in_type_tag)>;
-            using OutType = cuda_type_t<MLX_GET_TYPE(out_type_tag)>;
             using IdxT = std::conditional_t<large(), int64_t, int32_t>;
-            const InType* in_ptr = gpu_ptr<InType>(in) + offset_in;
-            OutType* out_ptr = gpu_ptr<OutType>(out) + offset_out;
-            int ndim = shape.size();
-            int work_per_thread = 1;
+
+            int work_per_thread = 8;
             auto dim0 = ndim > 0 ? shape.back() : 1;
             auto rest = out.size() / dim0;
-            if (dim0 >= 4) {
+            if (dim0 >= 4 && dim0 < 8) {
               work_per_thread = 4;
+            } else if (dim0 < 4) {
+              work_per_thread = 1;
             }
             dim0 = (dim0 + work_per_thread - 1) / work_per_thread;
             auto block_dims = get_block_dims(dim0, rest, 1);
@@ -110,7 +184,10 @@ void copy_general_input(
               dispatch_1_2_3(ndim, [&](auto dims_constant) {
                 auto kernel =
                     cu::copy_g_nd<InType, OutType, IdxT, dims_constant(), 1>;
-                if (work_per_thread == 4) {
+                if (work_per_thread == 8) {
+                  kernel =
+                      cu::copy_g_nd<InType, OutType, IdxT, dims_constant(), 8>;
+                } else if (work_per_thread == 4) {
                   kernel =
                       cu::copy_g_nd<InType, OutType, IdxT, dims_constant(), 4>;
                 }
@@ -127,7 +204,9 @@ void copy_general_input(
               });
             } else { // ndim >= 4
               auto kernel = cu::copy_g<InType, OutType, IdxT, 1>;
-              if (work_per_thread == 4) {
+              if (work_per_thread == 8) {
+                kernel = cu::copy_g<InType, OutType, IdxT, 8>;
+              } else if (work_per_thread == 4) {
                 kernel = cu::copy_g<InType, OutType, IdxT, 4>;
               }
               encoder.add_kernel_node(

mlx/backend/cuda/cublas_utils.cpp

@@ -0,0 +1,222 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/cublas_utils.h"
+#include "mlx/backend/cuda/cuda.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+namespace cublas_utils {
+
+namespace {
+
+struct CublasPreference {
+  CublasPreference(cu::Device& device) {
+    // The recommended cublas workspace size is 4 MiB for pre-Hopper and 32 MiB
+    // for Hopper+:
+    // https://docs.nvidia.com/cuda/cublas/#cublassetworkspace
+    uint64_t MiB = 1024 * 1024;
+    uint64_t workspace_size =
+        device.compute_capability_major() >= 9 ? 32 * MiB : 4 * MiB;
+
+    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceCreate(&pref_));
+    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceSetAttribute(
+        pref_,
+        CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
+        &workspace_size,
+        sizeof(uint64_t)));
+  }
+
+  ~CublasPreference() {
+    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceDestroy(pref_));
+  }
+
+  cublasLtMatmulPreference_t pref_{nullptr};
+};
+
+} // namespace
+
+cublasLtMatmulPreference_t get_preference(cu::Device& device) {
+  static CublasPreference pref(device);
+  return pref.pref_;
+}
+
+cublasLtMatrixLayout_t create_matrix_layout(
+    cudaDataType_t type,
+    uint64_t rows,
+    uint64_t cols,
+    bool transposed,
+    int64_t ld,
+    int32_t batch_count,
+    int64_t batch_stride) {
+  cublasLtMatrixLayout_t desc;
+  if (transposed) {
+    std::swap(rows, cols);
+  }
+  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutCreate(&desc, type, rows, cols, ld));
+  if (batch_count > 1) {
+    CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutSetAttribute(
+        desc,
+        CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT,
+        &batch_count,
+        sizeof(int32_t)));
+    CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutSetAttribute(
+        desc,
+        CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET,
+        &batch_stride,
+        sizeof(int64_t)));
+  }
+  return desc;
+}
+
+} // namespace cublas_utils
+
+CublasMatmulBase::~CublasMatmulBase() {
+  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(a_desc_));
+  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(b_desc_));
+  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(c_desc_));
+  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(out_desc_));
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescDestroy(matmul_desc_));
+}
+
+void CublasMatmulBase::init_base(
+    cu::Device& device,
+    cudaDataType_t scale_type,
+    cublasComputeType_t compute_type,
+    cudaDataType_t data_type,
+    cudaDataType_t output_type,
+    bool a_transposed,
+    uint64_t a_rows,
+    uint64_t a_cols,
+    int64_t lda,
+    bool b_transposed,
+    uint64_t b_rows,
+    uint64_t b_cols,
+    int64_t ldb,
+    int32_t batch_count,
+    int64_t a_batch_stride,
+    int64_t b_batch_stride) {
+  M_ = a_rows;
+  N_ = b_cols;
+  scale_type_ = scale_type;
+  handle_ = device.lt_handle();
+  pref_ = cublas_utils::get_preference(device);
+  heuristic_.state = CUBLAS_STATUS_NOT_INITIALIZED;
+
+  CHECK_CUBLAS_ERROR(
+      cublasLtMatmulDescCreate(&matmul_desc_, compute_type, scale_type));
+
+  int32_t pointer_mode = CUBLASLT_POINTER_MODE_HOST;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_POINTER_MODE,
+      &pointer_mode,
+      sizeof(int32_t)));
+
+  // In cublasLt matrices use column-major layout, while it is possible to use
+  // the CUBLASLT_ORDER_ROW option to switch to row-major layout, the bias
+  // epilogue does not work with the option. So instead we swap A and B to make
+  // cublasLt return the row-major result, which works because:
+  // - the data of a matrix in row-major layout is identical to its transpose in
+  //   column-major layout
+  // - C^T = (A @ B)^T = B^T @ A^T
+  cublasOperation_t a_op = b_transposed ? CUBLAS_OP_T : CUBLAS_OP_N;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_TRANSA,
+      &a_op,
+      sizeof(cublasOperation_t)));
+  cublasOperation_t b_op = a_transposed ? CUBLAS_OP_T : CUBLAS_OP_N;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_TRANSB,
+      &b_op,
+      sizeof(cublasOperation_t)));
+
+  a_desc_ = cublas_utils::create_matrix_layout(
+      data_type,
+      b_cols,
+      b_rows,
+      b_transposed,
+      ldb,
+      batch_count,
+      b_batch_stride);
+  b_desc_ = cublas_utils::create_matrix_layout(
+      data_type,
+      a_cols,
+      a_rows,
+      a_transposed,
+      lda,
+      batch_count,
+      a_batch_stride);
+  out_desc_ = cublas_utils::create_matrix_layout(
+      output_type, b_cols, a_rows, false, b_cols, batch_count, b_cols * a_rows);
+}
+
+void CublasMatmulBase::execute_matmul(
+    cu::CommandEncoder& encoder,
+    void* out,
+    const void* a,
+    const void* b,
+    const void* c,
+    const void* alpha_ptr,
+    const void* beta_ptr) {
+  if (heuristic_.state != CUBLAS_STATUS_SUCCESS) {
+    int ret = 0;
+    CHECK_CUBLAS_ERROR(cublasLtMatmulAlgoGetHeuristic(
+        handle_,
+        matmul_desc_,
+        a_desc_,
+        b_desc_,
+        c ? c_desc_ : out_desc_,
+        out_desc_,
+        pref_,
+        1,
+        &heuristic_,
+        &ret));
+    if (ret == 0) {
+      throw std::runtime_error("Can not find algorithm for matmul.");
+    }
+  }
+
+  void* workspace_ptr = allocate_workspace(encoder, heuristic_.workspaceSize);
+
+  // Execute matmul
+  auto capture = encoder.capture_context();
+  CHECK_CUBLAS_ERROR(cublasLtMatmul(
+      handle_,
+      matmul_desc_,
+      alpha_ptr,
+      b, // a and b are swapped for row-major layout
+      a_desc_,
+      a,
+      b_desc_,
+      beta_ptr,
+      c ? c : out,
+      c ? c_desc_ : out_desc_,
+      out,
+      out_desc_,
+      &heuristic_.algo,
+      workspace_ptr,
+      heuristic_.workspaceSize,
+      encoder.stream()));
+}
+
+void CublasMatmulBase::set_bias(
+    cu::CommandEncoder& encoder,
+    const array& bias) {
+  encoder.set_input_array(bias);
+  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_BIAS;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_EPILOGUE,
+      &epilogue,
+      sizeof(epilogue)));
+  auto* bias_ptr = gpu_ptr<void>(bias);
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_BIAS_POINTER,
+      &bias_ptr,
+      sizeof(bias_ptr)));
+}
+
+} // namespace mlx::core

mlx/backend/cuda/cublas_utils.h

@@ -0,0 +1,94 @@
+// Copyright © 2025 Apple Inc.
+#pragma once
+
+#include <cublasLt.h>
+#include "mlx/array.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/dtype_utils.h"
+
+namespace mlx::core {
+namespace cublas_utils {
+
+// Get the shared cublas preference for a device
+cublasLtMatmulPreference_t get_preference(cu::Device& device);
+
+cublasLtMatrixLayout_t create_matrix_layout(
+    cudaDataType_t type,
+    uint64_t rows,
+    uint64_t cols,
+    bool transposed,
+    int64_t ld,
+    int32_t batch_count,
+    int64_t batch_stride);
+
+inline cudaDataType_t dtype_to_cublas_type(Dtype dtype, std::string_view tag) {
+  switch (dtype) {
+    case float16:
+      return CUDA_R_16F;
+    case bfloat16:
+      return CUDA_R_16BF;
+    case float32:
+      return CUDA_R_32F;
+    case float64:
+      return CUDA_R_64F;
+    case complex64:
+      return CUDA_C_32F;
+    default:
+      throw std::runtime_error(fmt::format(
+          "Unsupported dtype in {}: {}.", tag, dtype_to_string(dtype)));
+  }
+}
+
+} // namespace cublas_utils
+
+class CublasMatmulBase {
+ public:
+  virtual ~CublasMatmulBase();
+
+  void set_bias(cu::CommandEncoder& encoder, const array& bias);
+
+ protected:
+  CublasMatmulBase() = default;
+
+  // Common member variables shared by all matmul types
+  uint64_t M_;
+  uint64_t N_;
+  cudaDataType_t scale_type_;
+  cublasLtMatmulPreference_t pref_{nullptr};
+  cublasLtHandle_t handle_{nullptr};
+  cublasLtMatmulDesc_t matmul_desc_{nullptr};
+  cublasLtMatrixLayout_t a_desc_{nullptr};
+  cublasLtMatrixLayout_t b_desc_{nullptr};
+  cublasLtMatrixLayout_t c_desc_{nullptr};
+  cublasLtMatrixLayout_t out_desc_{nullptr};
+  cublasLtMatmulHeuristicResult_t heuristic_;
+
+  void init_base(
+      cu::Device& device,
+      cudaDataType_t scale_type,
+      cublasComputeType_t compute_type,
+      cudaDataType_t data_type,
+      cudaDataType_t output_type,
+      bool a_transposed,
+      uint64_t a_rows,
+      uint64_t a_cols,
+      int64_t lda,
+      bool b_transposed,
+      uint64_t b_rows,
+      uint64_t b_cols,
+      int64_t ldb,
+      int32_t batch_count,
+      int64_t a_batch_stride,
+      int64_t b_batch_stride);
+
+  void execute_matmul(
+      cu::CommandEncoder& encoder,
+      void* out,
+      const void* a,
+      const void* b,
+      const void* c,
+      const void* alpha_ptr,
+      const void* beta_ptr);
+};
+
+} // namespace mlx::core

mlx/backend/cuda/cuda_utils.h

@@ -5,6 +5,7 @@
 #include <cublasLt.h>
 #include <cuda.h>
 #include <cuda_runtime.h>
+#include <cudnn.h>
 
 namespace mlx::core {
 
@@ -12,10 +13,12 @@ namespace mlx::core {
 void check_cublas_error(const char* name, cublasStatus_t err);
 void check_cuda_error(const char* name, cudaError_t err);
 void check_cuda_error(const char* name, CUresult err);
+void check_cudnn_error(const char* name, cudnnStatus_t err);
 
 // The macro version that prints the command that failed.
 #define CHECK_CUBLAS_ERROR(cmd) check_cublas_error(#cmd, (cmd))
 #define CHECK_CUDA_ERROR(cmd) check_cuda_error(#cmd, (cmd))
+#define CHECK_CUDNN_ERROR(cmd) check_cudnn_error(#cmd, (cmd))
 
 // Base class for RAII managed CUDA resources.
 template <typename Handle, cudaError_t (*Destroy)(Handle)>
@@ -29,6 +32,10 @@ class CudaHandle {
   }
 
   ~CudaHandle() {
+    // Skip if there was an error to avoid throwing in the destructors
+    if (cudaPeekAtLastError() != cudaSuccess) {
+      return;
+    }
     reset();
   }
 

mlx/backend/cuda/cudnn_utils.cpp

@@ -7,32 +7,26 @@ namespace mlx::core {
 
 namespace {
 
-// Create a cudnn tensor descriptor.
-template <typename Vec>
-inline cudnn_frontend::Tensor build_cudnn_tensor(
-    int64_t id,
-    const array& x,
-    const Vec& shape,
-    const Vec& strides) {
-  return cudnn_frontend::TensorBuilder()
-      .setDim(shape.size(), shape.data())
-      .setStrides(strides.size(), strides.data())
-      .setId(id)
-      .setAlignment(get_alignment(x))
-      .setDataType(dtype_to_cudnn_type(x.dtype()))
-      .build();
-}
+#define RETURN_IF_ERROR(cmd)          \
+  if (auto ret = cmd; ret.is_bad()) { \
+    return ret;                       \
+  }
 
 // In MLX a singleton dim (shape[dim] == 1) can have any stride, but in cuDNN
 // whether a tensor is contiguous is determined with:
 // shape[dim] == shape[dim + 1] * strides[dim + 1]
 // So a contiguous array with singleton dims in MLX may be mistakenly treated
 // as strided in cuDNN, and we work around it by normalizing the strides.
-Strides normalized_strides(const array& x) {
-  if (!x.flags().row_contiguous || x.ndim() < 2) {
-    return x.strides();
+std::vector<int64_t> normalized_strides(const array& x) {
+  std::vector<int64_t> strides(x.strides().begin(), x.strides().end());
+  if (std::all_of(
+          strides.begin(), strides.end(), [](int64_t s) { return s == 0; })) {
+    strides.back() = 1;
+    return strides;
+  }
+  if (!x.flags().row_contiguous || x.ndim() < 2) {
+    return strides;
   }
-  Strides strides = x.strides();
   for (int i = x.ndim() - 2; i >= 0; --i) {
     if (x.shape(i) == 1) {
       strides[i] = x.shape(i + 1) * strides[i + 1];
@@ -42,7 +36,9 @@ Strides normalized_strides(const array& x) {
 }
 
 // Return the shape and strides after transposing from NHWC to NCHW.
-auto nhwc_to_nchw(SmallVector<int64_t> shape, SmallVector<int64_t> strides) {
+inline auto nhwc_to_nchw(const array& x) {
+  auto shape = convert_vector<int64_t>(x.shape());
+  auto strides = normalized_strides(x);
   assert(shape.size() >= 3);
   shape.insert(shape.begin() + 1, shape.back());
   shape.erase(shape.end() - 1);
@@ -51,226 +47,87 @@ auto nhwc_to_nchw(SmallVector<int64_t> shape, SmallVector<int64_t> strides) {
   return std::make_tuple(std::move(shape), std::move(strides));
 }
 
-inline auto nhwc_to_nchw(const array& x) {
-  return nhwc_to_nchw(
-      convert_vector<int64_t>(x.shape()), normalized_strides(x));
-}
-
-// Return available engines for a |op_graph|.
-cudnn_frontend::EngineConfigList get_cudnn_engine_configs(
-    cudnnBackendDescriptorType_t backend_type,
-    Dtype dtype,
-    cudnn_frontend::OperationGraph& op_graph,
-    bool use_fallback = true) {
-  SmallVector<cudnn_frontend::GeneratorSource, 2> sources;
-  sources.push_back([](auto& op_graph) {
-    auto heuristics = cudnn_frontend::EngineHeuristicsBuilder()
-                          .setOperationGraph(op_graph)
-                          .setHeurMode(CUDNN_HEUR_MODE_A)
-                          .build();
-    return heuristics.getEngineConfig(heuristics.getEngineConfigCount());
-  });
-  if (use_fallback) {
-    sources.push_back([&backend_type](auto& op_graph) {
-      auto fallback = cudnn_frontend::EngineFallbackListBuilder()
-                          .setOperationGraph(op_graph)
-                          .setOperation(backend_type)
-                          .build();
-      return fallback.getFallbackList();
-    });
-  }
-
-  auto configs =
-      cudnn_frontend::EngineConfigGenerator(sources.size(), sources.data())
-          .generate_engine_config(op_graph);
-
-  cudnn_frontend::EngineConfigList filtered_configs;
-  cudnn_frontend::filter(configs, filtered_configs, [dtype](auto c) {
-    if (cudnn_frontend::hasNumericalNote<
-            CUDNN_NUMERICAL_NOTE_DOWN_CONVERT_INPUTS>(c)) {
-      return true;
-    }
-    if (cudnn_frontend::hasNumericalNote<CUDNN_NUMERICAL_NOTE_TENSOR_CORE>(c) &&
-        dtype == float32 && !env::enable_tf32()) {
-      return true;
-    }
-    return false;
-  });
-  return filtered_configs;
-}
-
-// Take |engine_configs| and |op_graph| and find a working execution plans
-// from them.
-std::optional<cudnn_frontend::ExecutionPlan>
-find_cudnn_plan_from_engine_configs(
-    cudnnHandle_t handle,
-    const cudnn_frontend::EngineConfigList& engine_configs,
-    const cudnn_frontend::OperationGraph& op_graph) {
-  auto op_graph_tag = op_graph.getTag();
-  for (const auto& config : engine_configs) {
-    try {
-      return cudnn_frontend::ExecutionPlanBuilder()
-          .setHandle(handle)
-          .setEngineConfig(config, op_graph_tag)
-          .build();
-    } catch (cudnn_frontend::cudnnException& error) {
-      if (error.getCudnnStatus() != CUDNN_STATUS_NOT_SUPPORTED) {
-        throw;
-      }
-    }
-  }
-  return std::nullopt;
-}
-
-// Prepare workspace and args to execute plan.
-template <typename F>
-bool prepare_cudnn_plan(
-    cu::CommandEncoder& encoder,
-    cudnn_frontend::ExecutionPlan& plan,
-    int num_args,
-    const int64_t* uids,
-    void** data_ptrs,
-    F&& execute) {
-  int workspace_size = plan.getWorkspaceSize();
-  void* workspace_ptr = nullptr;
-  if (workspace_size > 0) {
-    array workspace(
-        cu::malloc_async(workspace_size, encoder), {workspace_size}, uint8);
-    encoder.add_temporary(workspace);
-    workspace_ptr = gpu_ptr<void>(workspace);
-  }
-
-  auto args = cudnn_frontend::VariantPackBuilder()
-                  .setWorkspacePointer(workspace_ptr)
-                  .setDataPointers(num_args, data_ptrs)
-                  .setUids(num_args, uids)
-                  .build();
-
-  auto handle = encoder.device().cudnn_handle();
-  cudnnSetStream(handle, encoder.stream());
-
-  if (!execute(handle, plan.get_raw_desc(), args.get_raw_desc())) {
-    return false;
-  }
-
-  return true;
-}
-
 } // namespace
 
-cudnn_frontend::Tensor build_cudnn_tensor(int64_t id, const array& x) {
-  auto shape = convert_vector<int64_t>(x.shape());
-  return build_cudnn_tensor(id, x, shape, normalized_strides(x));
+fe::error_t DnnGraph::prepare() {
+  RETURN_IF_ERROR(validate());
+  try {
+    RETURN_IF_ERROR(build_operation_graph(handle_));
+  } catch (cudnn_frontend::cudnnException& error) {
+    // cuDNN bug: they did not catch all exceptions in the API.
+    return {fe::error_code_t::CUDNN_BACKEND_API_FAILED, error.what()};
+  }
+  RETURN_IF_ERROR(create_execution_plans({fe::HeurMode_t::A}));
+  return {};
 }
 
-cudnn_frontend::Tensor build_cudnn_tensor_nchw(int64_t id, const array& x) {
+fe::error_t DnnGraph::build() {
+  RETURN_IF_ERROR(check_support(handle_));
+  RETURN_IF_ERROR(build_plans(handle_));
+  return {};
+}
+
+fe::error_t DnnGraph::encode_graph(
+    cu::CommandEncoder& encoder,
+    std::unordered_map<int64_t, void*> variant_pack) {
+  cudnnSetStream(handle_, encoder.stream());
+  CudaGraph cuda_graph(encoder.device());
+  RETURN_IF_ERROR(populate_cuda_graph(
+      handle_, variant_pack, prepare_workspace(encoder), cuda_graph));
+  encoder.add_graph_node(cuda_graph);
+  return {};
+}
+
+fe::error_t DnnGraph::encode_capturing(
+    cu::CommandEncoder& encoder,
+    std::unordered_map<int64_t, void*> variant_pack) {
+  auto* workspace_ptr = prepare_workspace(encoder);
+  auto capture = encoder.capture_context();
+  cudnnSetStream(handle_, encoder.stream());
+  auto ret = execute(handle_, variant_pack, workspace_ptr);
+  if (ret.is_bad()) {
+    capture.discard = true;
+  }
+  return ret;
+}
+
+void* DnnGraph::prepare_workspace(cu::CommandEncoder& encoder) {
+  int64_t workspace_size = 0;
+  CHECK_CUDNN_FE_ERROR(get_workspace_size(workspace_size));
+  return allocate_workspace(encoder, workspace_size);
+}
+
+void DnnGraph::set_tensor_attrs(
+    std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
+    int64_t uid,
+    const array& x,
+    const std::vector<int64_t>& shape,
+    const std::vector<int64_t>& strides) {
+  tensor->set_uid(uid)
+      .set_alignment(get_alignment(x))
+      .set_data_type(dtype_to_cudnn_type(x.dtype()))
+      .set_dim(shape)
+      .set_stride(strides);
+}
+
+void DnnGraph::set_tensor_attrs(
+    std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
+    int64_t uid,
+    const array& x) {
+  set_tensor_attrs(
+      tensor,
+      uid,
+      x,
+      convert_vector<int64_t>(x.shape()),
+      normalized_strides(x));
+}
+
+void DnnGraph::set_tensor_attrs_nchw(
+    std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
+    int64_t uid,
+    const array& x) {
   auto [shape, strides] = nhwc_to_nchw(x);
-  return build_cudnn_tensor(id, x, shape, strides);
+  set_tensor_attrs(tensor, uid, x, shape, strides);
 }
 
-cudnn_frontend::Tensor build_cudnn_tensor_4d_nchw(int64_t id, const array& x) {
-  if (x.ndim() == 0) {
-    SmallVector<int64_t, 4> scalar_dims = {1, 1, 1, 1};
-    return build_cudnn_tensor(id, x, scalar_dims, scalar_dims);
-  }
-  if (x.ndim() == 1) {
-    int64_t s = x.shape(0);
-    SmallVector<int64_t, 4> shape = {1, x.shape(0), 1, 1};
-    SmallVector<int64_t, 4> strides = {s, 1, s, s};
-    return build_cudnn_tensor(id, x, shape, strides);
-  }
-  if (x.ndim() == 2) {
-    int64_t s =
-        x.flags().row_contiguous ? x.shape(1) * x.strides(1) : x.strides(0);
-    SmallVector<int64_t, 4> shape = {x.shape(0), x.shape(1), 1, 1};
-    SmallVector<int64_t, 4> strides = {s, x.strides(1), s, s};
-    return build_cudnn_tensor(id, x, shape, strides);
-  }
-  if (x.ndim() == 3 || x.ndim() == 4) {
-    return build_cudnn_tensor_nchw(id, x);
-  }
-  throw std::runtime_error(
-      fmt::format("Unsupported array with {} dims.", x.ndim()));
-}
-
-cudnn_frontend::Tensor build_cudnn_scalar_4d(int64_t id, Dtype dtype) {
-  SmallVector<int64_t, 4> scalar_dims = {1, 1, 1, 1};
-  return cudnn_frontend::TensorBuilder()
-      .setDim(scalar_dims.size(), scalar_dims.data())
-      .setStrides(scalar_dims.size(), scalar_dims.data())
-      .setId(id)
-      .setAlignment(16)
-      .setDataType(dtype_to_cudnn_type(dtype))
-      .setByValue(true)
-      .build();
-}
-
-std::optional<cudnn_frontend::ExecutionPlan> find_cudnn_plan_from_op_graph(
-    cudnnHandle_t handle,
-    cudnnBackendDescriptorType_t backend_type,
-    Dtype dtype,
-    cudnn_frontend::OperationGraph& op_graph) {
-  auto engine_configs = get_cudnn_engine_configs(backend_type, dtype, op_graph);
-  if (engine_configs.empty()) {
-    return std::nullopt;
-  }
-  return find_cudnn_plan_from_engine_configs(handle, engine_configs, op_graph);
-}
-
-bool encode_cudnn_plan_with_capturing(
-    cu::CommandEncoder& encoder,
-    cudnn_frontend::ExecutionPlan& plan,
-    int num_args,
-    const int64_t* uids,
-    void** data_ptrs) {
-  return prepare_cudnn_plan(
-      encoder,
-      plan,
-      num_args,
-      uids,
-      data_ptrs,
-      [&](auto handle, auto plan, auto args) {
-        auto capture = encoder.capture_context();
-        if (cudnnBackendExecute(handle, plan, args) != CUDNN_STATUS_SUCCESS) {
-          // Discard the captured graph when failed.
-          capture.discard = true;
-          return false;
-        }
-        return true;
-      });
-}
-
-#if CUDNN_VERSION >= 90500
-bool encode_cudnn_plan_with_graph_api(
-    cu::CommandEncoder& encoder,
-    cudnn_frontend::ExecutionPlan& plan,
-    CudaGraph& graph,
-    int num_args,
-    const int64_t* uids,
-    void** data_ptrs) {
-  return prepare_cudnn_plan(
-      encoder,
-      plan,
-      num_args,
-      uids,
-      data_ptrs,
-      [&](auto handle, auto plan, auto args) {
-        if (!graph) {
-          graph = CudaGraph(encoder.device());
-          if (cudnnBackendPopulateCudaGraph(handle, plan, args, graph) !=
-              CUDNN_STATUS_SUCCESS) {
-            return false;
-          }
-        } else {
-          if (cudnnBackendUpdateCudaGraph(handle, plan, args, graph) !=
-              CUDNN_STATUS_SUCCESS) {
-            return false;
-          }
-        }
-        encoder.add_graph_node(graph);
-        return true;
-      });
-}
-#endif
-
 } // namespace mlx::core

mlx/backend/cuda/cudnn_utils.h

@@ -2,25 +2,30 @@
 
 #pragma once
 
-#include "mlx/array.h"
-#include "mlx/backend/cuda/allocator.h"
 #include "mlx/backend/cuda/device/config.h"
 #include "mlx/backend/cuda/utils.h"
 #include "mlx/dtype_utils.h"
 
 #include <cudnn_frontend.h>
-#include <cudnn_frontend_find_plan.h>
 #include <fmt/format.h>
 
-#include <algorithm>
-#include <array>
-
 namespace mlx::core {
 
 namespace cu {
 class CommandEncoder;
 }
 
+namespace fe = cudnn_frontend;
+
+#define CHECK_CUDNN_FE_ERROR(cmd)                                    \
+  do {                                                               \
+    auto error = cmd;                                                \
+    if (!error.is_good()) {                                          \
+      throw std::runtime_error(                                      \
+          fmt::format("{} failed: {}.", #cmd, error.get_message())); \
+    }                                                                \
+  } while (0)
+
 // Return pointer alignment of |x|'s data.
 inline uint8_t get_alignment(const array& x) {
   uint8_t alignment = 1;
@@ -35,8 +40,31 @@ inline uint8_t get_alignment(const array& x) {
 
 // Convert the type of elements in |vec| to |T|.
 template <typename T, typename Vec>
-inline SmallVector<T> convert_vector(const Vec& vec) {
-  return SmallVector<T>(vec.begin(), vec.end());
+inline std::vector<T> convert_vector(const Vec& vec) {
+  return std::vector<T>(vec.begin(), vec.end());
+}
+
+// Map dtype to cudnn data type.
+inline fe::DataType_t dtype_to_cudnn_type(Dtype dtype) {
+  switch (dtype) {
+    case int8:
+      return fe::DataType_t::INT8;
+    case int32:
+      return fe::DataType_t::INT32;
+    case uint8:
+      return fe::DataType_t::UINT8;
+    case float16:
+      return fe::DataType_t::HALF;
+    case bfloat16:
+      return fe::DataType_t::BFLOAT16;
+    case float32:
+      return fe::DataType_t::FLOAT;
+    case float64:
+      return fe::DataType_t::DOUBLE;
+    default:
+      throw std::runtime_error(fmt::format(
+          "Unsupported dtype in cuDNN: {}.", dtype_to_string(dtype)));
+  }
 }
 
 // Return an array that can be used as map key for |vec| with size <= MAX_NDIM.
@@ -55,111 +83,89 @@ inline std::array<T, NDIM> vector_key(const Vec<T>& vec) {
   return result;
 }
 
-// Helpers used by get_data_ptrs to get pointers.
-inline void* get_data_ptr(const array& arr) {
-  return const_cast<void*>(gpu_ptr<void>(arr));
-}
-
-template <typename T, typename = std::enable_if_t<std::is_scalar_v<T>>>
-inline void* get_data_ptr(T& scalar) {
-  return &scalar;
-}
-
-// Return an array filled with data pointers of args.
-template <typename... Args>
-inline std::array<void*, sizeof...(Args)> get_data_ptrs(Args&... args) {
-  return {get_data_ptr(args)...};
-}
-
-// Map dtype to cudnn data type.
-inline cudnnDataType_t dtype_to_cudnn_type(Dtype dtype) {
-  switch (dtype) {
-    case int8:
-      return CUDNN_DATA_INT8;
-    case int32:
-      return CUDNN_DATA_INT32;
-    case uint8:
-      return CUDNN_DATA_UINT8;
-    case float16:
-      return CUDNN_DATA_HALF;
-    case bfloat16:
-      return CUDNN_DATA_BFLOAT16;
-    case float32:
-      return CUDNN_DATA_FLOAT;
-    case float64:
-      return CUDNN_DATA_DOUBLE;
-    default:
-      throw std::runtime_error(fmt::format(
-          "Unsupported dtype in Convolution: {}.", dtype_to_string(dtype)));
+// Extends cuDNN graph with helpers.
+class DnnGraph : public fe::graph::Graph {
+ public:
+  DnnGraph(cudnnHandle_t handle, Dtype io_dtype, Dtype compute_dtype = float32)
+      : handle_(handle) {
+    set_io_data_type(dtype_to_cudnn_type(io_dtype));
+    set_intermediate_data_type(dtype_to_cudnn_type(compute_dtype));
+    set_compute_data_type(dtype_to_cudnn_type(compute_dtype));
   }
-}
 
-// Create a tensor descriptor from |x|.
-cudnn_frontend::Tensor build_cudnn_tensor(int64_t id, const array& x);
+  // Create a cuDNN tensor description from MLX array |x|.
+  auto& tensor(
+      std::shared_ptr<fe::graph::Tensor_attributes>& attrs,
+      int64_t uid,
+      const array& x) {
+    set_tensor_attrs(attrs, uid, x);
+    return attrs;
+  }
+  auto tensor(const char* name, int64_t uid, const array& x) {
+    auto attrs = Graph::tensor(fe::graph::Tensor_attributes().set_name(name));
+    tensor(attrs, uid, x);
+    return attrs;
+  }
 
-// Create a tensor descriptor from |x|, and transpose from NHWC to NCHW.
-cudnn_frontend::Tensor build_cudnn_tensor_nchw(int64_t id, const array& x);
+  // Create a cuDNN tensor description from MLX array |x|, and transpose it from
+  // NHWC layout to NCHW.
+  auto& tensor_nchw(
+      std::shared_ptr<fe::graph::Tensor_attributes>& attrs,
+      int64_t uid,
+      const array& x) {
+    set_tensor_attrs_nchw(attrs, uid, x);
+    return attrs;
+  }
+  auto tensor_nchw(const char* name, int64_t uid, const array& x) {
+    auto attrs = Graph::tensor(fe::graph::Tensor_attributes().set_name(name));
+    tensor_nchw(attrs, uid, x);
+    return attrs;
+  }
 
-// Create a tensor descriptor from |x|, make sure it is 4D, and transpose it
-// from NHWC to NCHW.
-cudnn_frontend::Tensor build_cudnn_tensor_4d_nchw(int64_t id, const array& x);
+  // Create a cuDNN tensor for scalar.
+  auto scalar(const char* name, int64_t uid, Dtype dtype) {
+    return Graph::tensor(fe::graph::Tensor_attributes()
+                             .set_name(name)
+                             .set_uid(uid)
+                             .set_dim({1, 1, 1, 1})
+                             .set_stride({1, 1, 1, 1})
+                             .set_is_pass_by_value(true)
+                             .set_data_type(dtype_to_cudnn_type(dtype)));
+  }
 
-// Create a 4D scalar tensor descriptor, which is passed by value.
-cudnn_frontend::Tensor build_cudnn_scalar_4d(int64_t id, Dtype dtype);
+  // Call this before setting notes.
+  fe::error_t prepare();
+  // Call this after setting notes.
+  fe::error_t build();
 
-// Find a working plan for |op_graph|.
-std::optional<cudnn_frontend::ExecutionPlan> find_cudnn_plan_from_op_graph(
-    cudnnHandle_t handle,
-    cudnnBackendDescriptorType_t backend_type,
-    Dtype dtype,
-    cudnn_frontend::OperationGraph& op_graph);
+  // Add cuDNN graph to CUDA graph, using native CUDA graph API.
+  fe::error_t encode_graph(
+      cu::CommandEncoder& encoder,
+      std::unordered_map<int64_t, void*> variant_pack);
+  // Add cuDNN graph to CUDA graph, using stream capture.
+  fe::error_t encode_capturing(
+      cu::CommandEncoder& encoder,
+      std::unordered_map<int64_t, void*> variant_pack);
 
-// Encode the plan to command buffer by capturing.
-bool encode_cudnn_plan_with_capturing(
-    cu::CommandEncoder& encoder,
-    cudnn_frontend::ExecutionPlan& plan,
-    int num_args,
-    const int64_t* uids,
-    void** data_ptrs);
+ private:
+  void* prepare_workspace(cu::CommandEncoder& encoder);
 
-#if CUDNN_VERSION >= 90500
-// Encode the plan to command buffer by using native graph api of cudnn. If the
-// |graph| is empty it will be populated, otherwise it will be updated.
-bool encode_cudnn_plan_with_graph_api(
-    cu::CommandEncoder& encoder,
-    cudnn_frontend::ExecutionPlan& plan,
-    CudaGraph& graph,
-    int num_args,
-    const int64_t* uids,
-    void** data_ptrs);
-#endif
+  void set_tensor_attrs(
+      std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
+      int64_t uid,
+      const array& x,
+      const std::vector<int64_t>& shape,
+      const std::vector<int64_t>& strides);
+  void set_tensor_attrs(
+      std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
+      int64_t uid,
+      const array& x);
+  void set_tensor_attrs_nchw(
+      std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
+      int64_t uid,
+      const array& x);
 
-// Helpers to make calls like encode_cudnn_plan(..., {'x', 'y', 'z'}, x, y, z).
-template <typename... Args>
-bool encode_cudnn_plan(
-    cu::CommandEncoder& encoder,
-    cudnn_frontend::ExecutionPlan& plan,
-    std::initializer_list<int64_t> uids,
-    Args&... args) {
-  assert(uids.size() == sizeof...(args));
-  auto data_ptrs = get_data_ptrs(args...);
-  return encode_cudnn_plan_with_capturing(
-      encoder, plan, uids.size(), uids.begin(), data_ptrs.data());
-}
-
-#if CUDNN_VERSION >= 90500
-template <typename... Args>
-bool encode_cudnn_plan(
-    cu::CommandEncoder& encoder,
-    cudnn_frontend::ExecutionPlan& plan,
-    CudaGraph& graph,
-    std::initializer_list<int64_t> uids,
-    Args&... args) {
-  assert(uids.size() == sizeof...(args));
-  auto data_ptrs = get_data_ptrs(args...);
-  return encode_cudnn_plan_with_graph_api(
-      encoder, plan, graph, uids.size(), uids.begin(), data_ptrs.data());
-}
-#endif
+  cudnnHandle_t handle_;
+};
 
 } // namespace mlx::core

mlx/backend/cuda/detect_cuda_arch.sh

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

mlx/backend/cuda/device.cpp

@@ -14,20 +14,20 @@ namespace mlx::core::cu {
 
 namespace {
 
-#define CHECK_CUDNN_ERROR(cmd) check_cudnn_error(#cmd, (cmd))
-
-void check_cudnn_error(const char* name, cudnnStatus_t err) {
-  if (err != CUDNN_STATUS_SUCCESS) {
-    throw std::runtime_error(
-        fmt::format("{} failed: {}.", name, cudnnGetErrorString(err)));
-  }
+bool use_cuda_graphs() {
+  static bool use_graphs = env::get_var("MLX_USE_CUDA_GRAPHS", true);
+  return use_graphs;
 }
 
-bool use_cuda_graphs() {
-  static bool use_graphs = []() {
-    return env::get_var("MLX_USE_CUDA_GRAPHS", true);
+const char* save_cuda_graphs_dot_file() {
+  static const char* filename = []() -> const char* {
+    const char* env = std::getenv("MLX_SAVE_CUDA_GRAPHS_DOT_FILE");
+    if (env && std::strlen(env) == 0) {
+      return nullptr;
+    }
+    return env;
   }();
-  return use_graphs;
+  return filename;
 }
 
 } // namespace
@@ -87,7 +87,7 @@ CommandEncoder::CaptureContext::CaptureContext(CommandEncoder& enc) : enc(enc) {
     return;
   }
   CHECK_CUDA_ERROR(
-      cudaStreamBeginCapture(enc.stream(), cudaStreamCaptureModeGlobal));
+      cudaStreamBeginCapture(enc.stream(), cudaStreamCaptureModeThreadLocal));
 }
 
 CommandEncoder::CaptureContext::~CaptureContext() {
@@ -115,18 +115,17 @@ CommandEncoder::ConcurrentContext::~ConcurrentContext() {
   }
 
   // Use an empty graph node for synchronization
-  CommandEncoder::GraphNode empty{NULL, 'E', std::to_string(enc.node_count_++)};
-  enc.empty_node_count_++;
+  CommandEncoder::GraphNode empty{NULL, "E", std::to_string(enc.node_count_++)};
   CHECK_CUDA_ERROR(cudaGraphAddEmptyNode(&empty.node, enc.graph_, NULL, 0));
 
   // Insert the concurrent -> empty node dependencies
   for (auto& from : enc.concurrent_nodes_) {
     enc.from_nodes_.push_back(from.node);
     enc.to_nodes_.push_back(empty.node);
-    enc.graph_key_ += from.id;
-    enc.graph_key_ += from.node_type;
-    enc.graph_key_ += empty.id;
-    enc.graph_key_ += empty.node_type;
+    enc.graph_deps_key_ += from.id;
+    enc.graph_deps_key_ += "-";
+    enc.graph_deps_key_ += empty.id;
+    enc.graph_deps_key_ += "-";
   }
 
   // Insert the input -> concurrent node dependencies without updating output
@@ -141,9 +140,6 @@ CommandEncoder::ConcurrentContext::~ConcurrentContext() {
 }
 
 void CommandEncoder::insert_graph_dependencies(GraphNode node) {
-  if (node.node_type == 'G') {
-    graph_node_count_++;
-  }
   node.id = std::to_string(node_count_++);
   if (in_concurrent_) {
     concurrent_nodes_.push_back(std::move(node));
@@ -155,6 +151,10 @@ void CommandEncoder::insert_graph_dependencies(GraphNode node) {
 }
 
 void CommandEncoder::insert_graph_dependencies(std::vector<GraphNode> nodes) {
+  for (auto& node : nodes) {
+    graph_nodes_key_ += node.node_type;
+    graph_nodes_key_ += "-";
+  }
   std::vector<GraphNode> deps;
   {
     // Dependencies must be added in the same order to produce a consistent
@@ -182,10 +182,10 @@ void CommandEncoder::insert_graph_dependencies(std::vector<GraphNode> nodes) {
     for (auto& to : nodes) {
       from_nodes_.push_back(from.node);
       to_nodes_.push_back(to.node);
-      graph_key_ += from.id;
-      graph_key_ += from.node_type;
-      graph_key_ += to.id;
-      graph_key_ += to.node_type;
+      graph_deps_key_ += from.id;
+      graph_deps_key_ += "-";
+      graph_deps_key_ += to.id;
+      graph_deps_key_ += "-";
     }
   }
 }
@@ -309,13 +309,76 @@ void CommandEncoder::add_kernel_node(
 void CommandEncoder::add_kernel_node(const cudaKernelNodeParams& params) {
   cudaGraphNode_t node;
   CHECK_CUDA_ERROR(cudaGraphAddKernelNode(&node, graph_, NULL, 0, &params));
-  insert_graph_dependencies(GraphNode{node, 'K'});
+  insert_graph_dependencies(GraphNode{node, "K"});
 }
 
 void CommandEncoder::add_kernel_node(const CUDA_KERNEL_NODE_PARAMS& params) {
   CUgraphNode node;
   CHECK_CUDA_ERROR(cuGraphAddKernelNode(&node, graph_, NULL, 0, &params));
-  insert_graph_dependencies(GraphNode{node, 'K'});
+  insert_graph_dependencies(GraphNode{node, "K"});
+}
+
+std::pair<std::string, bool> subgraph_to_key(cudaGraph_t graph) {
+  // Constructs a key representing the nodes of a sub-graph.
+  // Also checks if the sub-graph is updatable as CUDA graphs do not get
+  // updated correctly if a kernel node getting updated has a different cluster
+  // shape than the node it's being updated with.
+  std::string key = "(";
+  size_t num_nodes = 0;
+  CHECK_CUDA_ERROR(cudaGraphGetNodes(graph, nullptr, &num_nodes));
+  if (num_nodes == 0) {
+    return {key + ")", true};
+  }
+  bool is_updatable = true;
+  std::vector<cudaGraphNode_t> nodes(num_nodes);
+  CHECK_CUDA_ERROR(cudaGraphGetNodes(graph, nodes.data(), &num_nodes));
+  for (const auto& node : nodes) {
+    if (!is_updatable) {
+      break;
+    }
+    cudaGraphNodeType type;
+    CHECK_CUDA_ERROR(cudaGraphNodeGetType(node, &type));
+    switch (type) {
+      case cudaGraphNodeTypeGraph: {
+        // Try to be updatable for a structure like graph -> graph -> kernel
+        cudaGraph_t child;
+        CHECK_CUDA_ERROR(cudaGraphChildGraphNodeGetGraph(node, &child));
+        auto [subkey, sub_is_updatable] = subgraph_to_key(child);
+        is_updatable &= sub_is_updatable;
+        key += subkey;
+        break;
+      }
+      case cudaGraphNodeTypeHost:
+        key += "H";
+        break;
+      case cudaGraphNodeTypeMemset:
+        key += "M";
+        break;
+      case cudaGraphNodeTypeKernel: {
+        cudaLaunchAttributeValue cluster_dim;
+        CHECK_CUDA_ERROR(cudaGraphKernelNodeGetAttribute(
+            node, cudaLaunchAttributeClusterDimension, &cluster_dim));
+        // Only allow dim.x to be greater than 1
+        if (cluster_dim.clusterDim.y > 1 || cluster_dim.clusterDim.z > 1) {
+          is_updatable = false;
+        } else {
+          key += "K";
+          key += std::to_string(cluster_dim.clusterDim.x);
+        }
+        break;
+      }
+      case cudaGraphNodeTypeWaitEvent:
+        key += "W";
+        break;
+      case cudaGraphNodeTypeEventRecord:
+        key += "R";
+        break;
+      default:
+        is_updatable = false;
+    }
+  }
+  key += ")";
+  return {key, is_updatable};
 }
 
 void CommandEncoder::add_graph_node(cudaGraph_t child) {
@@ -328,8 +391,10 @@ void CommandEncoder::add_graph_node(cudaGraph_t child) {
     return;
   }
   cudaGraphNode_t node;
+  auto [sub_graph_key, is_updatable] = subgraph_to_key(child);
+  is_graph_updatable_ &= is_updatable;
   CHECK_CUDA_ERROR(cudaGraphAddChildGraphNode(&node, graph_, NULL, 0, child));
-  insert_graph_dependencies(GraphNode{node, 'G'});
+  insert_graph_dependencies(GraphNode{node, sub_graph_key});
 }
 
 bool CommandEncoder::needs_commit() {
@@ -354,44 +419,53 @@ void CommandEncoder::commit() {
           from_nodes_.size()));
     }
 
-    graph_key_ += ".";
-    graph_key_ += std::to_string(node_count_);
-    graph_key_ += ".";
-    graph_key_ += std::to_string(graph_node_count_);
-    graph_key_ += ".";
-    graph_key_ += std::to_string(empty_node_count_);
-
-    CudaGraphExec& graph_exec = graph_cache_[graph_key_];
-
-    if (graph_exec != nullptr) {
-      cudaGraphExecUpdateResult update_result;
-#if CUDART_VERSION >= 12000
-      cudaGraphExecUpdateResultInfo info;
-      cudaGraphExecUpdate(graph_exec, graph_, &info);
-      update_result = info.result;
-#else
-      cudaGraphNode_t error_node;
-      cudaGraphExecUpdate(graph_exec, graph_, &error_node, &update_result);
-#endif // CUDART_VERSION >= 12000
-      if (update_result != cudaGraphExecUpdateSuccess) {
-        cudaGetLastError(); // reset error
-        graph_exec.reset();
-      }
-    }
-    if (graph_exec == nullptr) {
-      graph_exec.instantiate(graph_);
-    }
     device_.make_current();
-    CHECK_CUDA_ERROR(cudaGraphLaunch(graph_exec, stream_));
+
+    if (!is_graph_updatable_) {
+      CudaGraphExec graph_exec;
+      graph_exec.instantiate(graph_);
+      CHECK_CUDA_ERROR(cudaGraphLaunch(graph_exec, stream_));
+    } else {
+      auto graph_key = graph_nodes_key_ + ":" + graph_deps_key_;
+      auto& graph_exec = graph_cache_[graph_key];
+
+      if (graph_exec != nullptr) {
+        cudaGraphExecUpdateResult update_result;
+#if CUDART_VERSION >= 12000
+        cudaGraphExecUpdateResultInfo info;
+        cudaGraphExecUpdate(graph_exec, graph_, &info);
+        update_result = info.result;
+#else
+        cudaGraphNode_t error_node;
+        cudaGraphExecUpdate(graph_exec, graph_, &error_node, &update_result);
+#endif // CUDART_VERSION >= 12000
+        if (update_result != cudaGraphExecUpdateSuccess) {
+          cudaGetLastError(); // reset error
+          graph_exec.reset();
+        }
+      }
+      if (graph_exec == nullptr) {
+        graph_exec.instantiate(graph_);
+      }
+
+      CHECK_CUDA_ERROR(cudaGraphLaunch(graph_exec, stream_));
+    }
+
+    // Save cuda graph to dot file
+    if (const char* filename = save_cuda_graphs_dot_file(); filename) {
+      static int count = 0;
+      auto path = fmt::format("{}_{}.dot", filename, ++count);
+      CHECK_CUDA_ERROR(cudaGraphDebugDotPrint(graph_, path.c_str(), 0));
+    }
 
     // Reset state
-    graph_node_count_ = 0;
-    empty_node_count_ = 0;
     from_nodes_.clear();
     to_nodes_.clear();
-    graph_key_.clear();
+    graph_deps_key_.clear();
+    graph_nodes_key_.clear();
     node_map_.clear();
     graph_ = CudaGraph(device_);
+    is_graph_updatable_ = true;
   }
 
   // Put completion handlers in a batch.

mlx/backend/cuda/device.h

@@ -106,8 +106,9 @@ class CommandEncoder {
     cudaGraphNode_t node;
     // K = kernel
     // E = empty
-    // G = subgraph
-    char node_type;
+    // () = subgraph (with metadata)
+    // Symbols ':', '-' are reserved as separators
+    std::string node_type;
     std::string id;
   };
 
@@ -119,12 +120,11 @@ class CommandEncoder {
   CudaGraph graph_;
   Worker worker_;
   char node_count_{0};
-  char graph_node_count_{0};
-  char empty_node_count_{0};
   bool in_concurrent_{false};
   std::vector<cudaGraphNode_t> from_nodes_;
   std::vector<cudaGraphNode_t> to_nodes_;
-  std::string graph_key_;
+  std::string graph_nodes_key_;
+  std::string graph_deps_key_;
   std::vector<GraphNode> concurrent_nodes_;
   std::vector<std::shared_ptr<array::Data>> temporaries_;
   LRUCache<std::string, CudaGraphExec> graph_cache_;
@@ -132,6 +132,7 @@ class CommandEncoder {
   std::vector<std::uintptr_t> active_outputs_;
   std::unordered_map<std::uintptr_t, GraphNode> node_map_;
   size_t bytes_in_graph_{0};
+  bool is_graph_updatable_{true};
   int max_ops_per_graph_;
   int max_mb_per_graph_;
 };

mlx/backend/cuda/event.cu

@@ -305,6 +305,7 @@ void Event::wait() {
   } else {
     event->atomic->wait(value());
   }
+  CHECK_CUDA_ERROR(cudaPeekAtLastError());
 }
 
 void Event::wait(Stream s) {

mlx/backend/cuda/gemms/cublas_gemm.cpp

@@ -1,6 +1,7 @@
 // Copyright © 2025 Apple Inc.
 
 #include "mlx/backend/cuda/gemms/cublas_gemm.h"
+#include "mlx/backend/cuda/cublas_utils.h"
 #include "mlx/backend/cuda/device.h"
 #include "mlx/dtype_utils.h"
 #include "mlx/utils.h"
@@ -11,35 +12,6 @@ namespace mlx::core {
 
 namespace {
 
-struct CublasPreference {
-  CublasPreference(cu::Device& device) {
-    // The recommended cublas workspace size is 4 MiB for pre-Hopper and 32 MiB
-    // for Hopper+:
-    // https://docs.nvidia.com/cuda/cublas/#cublassetworkspace
-    uint64_t MiB = 1024 * 1024;
-    uint64_t workspace_size =
-        device.compute_capability_major() >= 9 ? 32 * MiB : 4 * MiB;
-
-    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceCreate(&pref_));
-    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceSetAttribute(
-        pref_,
-        CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
-        &workspace_size,
-        sizeof(uint64_t)));
-  }
-
-  ~CublasPreference() {
-    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceDestroy(pref_));
-  }
-
-  cublasLtMatmulPreference_t pref_{nullptr};
-};
-
-cublasLtMatmulPreference_t cublas_preference(cu::Device& device) {
-  static CublasPreference pref(device);
-  return pref.pref_;
-}
-
 cublasComputeType_t dtype_to_compute_type(Dtype dtype) {
   switch (dtype) {
     case float16:
@@ -60,52 +32,6 @@ cublasComputeType_t dtype_to_compute_type(Dtype dtype) {
   }
 }
 
-cudaDataType_t dtype_to_cublas_type(Dtype dtype) {
-  switch (dtype) {
-    case float16:
-      return CUDA_R_16F;
-    case bfloat16:
-      return CUDA_R_16BF;
-    case float32:
-      return CUDA_R_32F;
-    case float64:
-      return CUDA_R_64F;
-    case complex64:
-      return CUDA_C_32F;
-    default:
-      throw std::runtime_error(fmt::format(
-          "Unsupported dtype in CublasGemm: {}.", dtype_to_string(dtype)));
-  }
-}
-
-cublasLtMatrixLayout_t create_matrix_layout(
-    cudaDataType_t type,
-    uint64_t rows,
-    uint64_t cols,
-    bool transposed,
-    int64_t ld,
-    int32_t batch_count,
-    int64_t batch_stride) {
-  cublasLtMatrixLayout_t desc;
-  if (transposed) {
-    std::swap(rows, cols);
-  }
-  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutCreate(&desc, type, rows, cols, ld));
-  if (batch_count > 1) {
-    CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutSetAttribute(
-        desc,
-        CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT,
-        &batch_count,
-        sizeof(int32_t)));
-    CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutSetAttribute(
-        desc,
-        CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET,
-        &batch_stride,
-        sizeof(int64_t)));
-  }
-  return desc;
-}
-
 } // namespace
 
 CublasGemm::CublasGemm(
@@ -121,54 +47,31 @@ CublasGemm::CublasGemm(
     int64_t ldb,
     int32_t batch_count,
     int64_t a_batch_stride,
-    int64_t b_batch_stride)
-    : handle_(device.lt_handle()),
-      pref_(cublas_preference(device)),
-      M_(a_rows),
-      N_(b_cols) {
-  heuristic_.state = CUBLAS_STATUS_NOT_INITIALIZED;
-
-  scale_type_ = dtype_to_cublas_type(dtype);
+    int64_t b_batch_stride) {
+  scale_type_ = cublas_utils::dtype_to_cublas_type(dtype, "CublasGemm");
   if (dtype == bfloat16 || dtype == float16) {
     scale_type_ = CUDA_R_32F;
   }
+  cudaDataType_t cublas_dtype =
+      cublas_utils::dtype_to_cublas_type(dtype, "CublasGemm");
 
-  CHECK_CUBLAS_ERROR(cublasLtMatmulDescCreate(
-      &matmul_desc_, dtype_to_compute_type(dtype), scale_type_));
-  int32_t pointer_mode = CUBLASLT_POINTER_MODE_HOST;
-  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
-      matmul_desc_,
-      CUBLASLT_MATMUL_DESC_POINTER_MODE,
-      &pointer_mode,
-      sizeof(int32_t)));
-
-  // In cublasLt matrices use column-major layout, while it is possible to use
-  // the CUBLASLT_ORDER_ROW option to switch to row-major layout, the bias
-  // epilogue does not work with the option. So instead we swap A and B to make
-  // cublasLt return the row-major result, which works because:
-  // - the data of a matrix in row-major layout is identical to its transpose in
-  //   column-major layout
-  // - C^T = (A @ B)^T = B^T @ A^T
-  cublasOperation_t a_op = b_transposed ? CUBLAS_OP_T : CUBLAS_OP_N;
-  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
-      matmul_desc_,
-      CUBLASLT_MATMUL_DESC_TRANSA,
-      &a_op,
-      sizeof(cublasOperation_t)));
-  cublasOperation_t b_op = a_transposed ? CUBLAS_OP_T : CUBLAS_OP_N;
-  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
-      matmul_desc_,
-      CUBLASLT_MATMUL_DESC_TRANSB,
-      &b_op,
-      sizeof(cublasOperation_t)));
-
-  auto type = dtype_to_cublas_type(dtype);
-  a_desc_ = create_matrix_layout(
-      type, b_cols, b_rows, b_transposed, ldb, batch_count, b_batch_stride);
-  b_desc_ = create_matrix_layout(
-      type, a_cols, a_rows, a_transposed, lda, batch_count, a_batch_stride);
-  out_desc_ = create_matrix_layout(
-      type, b_cols, a_rows, false, b_cols, batch_count, a_rows * b_cols);
+  init_base(
+      device,
+      scale_type_,
+      dtype_to_compute_type(dtype),
+      cublas_dtype,
+      cublas_dtype,
+      a_transposed,
+      a_rows,
+      a_cols,
+      lda,
+      b_transposed,
+      b_rows,
+      b_cols,
+      ldb,
+      batch_count,
+      a_batch_stride,
+      b_batch_stride);
 }
 
 CublasGemm::CublasGemm(
@@ -201,19 +104,11 @@ CublasGemm::CublasGemm(
           batch_count,
           a_batch_stride,
           b_batch_stride) {
-  auto type = dtype_to_cublas_type(dtype);
-  c_desc_ = create_matrix_layout(
+  auto type = cublas_utils::dtype_to_cublas_type(dtype, "CublasGemm");
+  c_desc_ = cublas_utils::create_matrix_layout(
       type, b_cols, a_rows, false, ldc, batch_count, c_batch_stride);
 }
 
-CublasGemm::~CublasGemm() {
-  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(a_desc_));
-  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(b_desc_));
-  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(c_desc_));
-  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(out_desc_));
-  CHECK_CUBLAS_ERROR(cublasLtMatmulDescDestroy(matmul_desc_));
-}
-
 void CublasGemm::set_out(
     Dtype dtype,
     bool transposed,
@@ -223,8 +118,8 @@ void CublasGemm::set_out(
     int32_t batch_count,
     int64_t batch_stride) {
   CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(out_desc_));
-  out_desc_ = create_matrix_layout(
-      dtype_to_cublas_type(dtype),
+  out_desc_ = cublas_utils::create_matrix_layout(
+      cublas_utils::dtype_to_cublas_type(dtype, "CublasGemm"),
       cols,
       rows,
       transposed,
@@ -233,22 +128,6 @@ void CublasGemm::set_out(
       batch_stride);
 }
 
-void CublasGemm::set_bias(cu::CommandEncoder& encoder, const array& bias) {
-  encoder.set_input_array(bias);
-  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_BIAS;
-  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
-      matmul_desc_,
-      CUBLASLT_MATMUL_DESC_EPILOGUE,
-      &epilogue,
-      sizeof(epilogue)));
-  auto* bias_ptr = gpu_ptr<void>(bias);
-  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
-      matmul_desc_,
-      CUBLASLT_MATMUL_DESC_BIAS_POINTER,
-      &bias_ptr,
-      sizeof(bias_ptr)));
-}
-
 void CublasGemm::run(
     cu::CommandEncoder& encoder,
     array& out,
@@ -337,24 +216,6 @@ void CublasGemm::execute(
     const void* c,
     float alpha /* = 1 */,
     float beta /* = 0 */) {
-  if (heuristic_.state != CUBLAS_STATUS_SUCCESS) {
-    int ret = 0;
-    CHECK_CUBLAS_ERROR(cublasLtMatmulAlgoGetHeuristic(
-        handle_,
-        matmul_desc_,
-        a_desc_,
-        b_desc_,
-        c ? c_desc_ : out_desc_,
-        out_desc_,
-        pref_,
-        1,
-        &heuristic_,
-        &ret));
-    if (ret == 0) {
-      throw std::runtime_error("Can not find algorithm for matmul.");
-    }
-  }
-
   const void* alpha_ptr = &alpha;
   const void* beta_ptr = &beta;
   complex64_t alpha_c, beta_c;
@@ -365,36 +226,7 @@ void CublasGemm::execute(
     beta_ptr = &beta_c;
   }
 
-  void* workspace_ptr = nullptr;
-  if (heuristic_.workspaceSize > 0) {
-    // Ensure workspace is 256-byte aligned
-    int nbytes = cuda::ceil_div(heuristic_.workspaceSize, 256) * 256;
-    array workspace(
-        cu::malloc_async(nbytes, encoder),
-        {static_cast<int>(heuristic_.workspaceSize)},
-        int8);
-    encoder.add_temporary(workspace);
-    workspace_ptr = gpu_ptr<void>(workspace);
-  }
-
-  auto capture = encoder.capture_context();
-  CHECK_CUBLAS_ERROR(cublasLtMatmul(
-      handle_,
-      matmul_desc_,
-      alpha_ptr,
-      b, // a and b are swapped
-      a_desc_,
-      a,
-      b_desc_,
-      beta_ptr,
-      c ? c : out,
-      c ? c_desc_ : out_desc_,
-      out,
-      out_desc_,
-      &heuristic_.algo,
-      workspace_ptr,
-      heuristic_.workspaceSize,
-      encoder.stream()));
+  execute_matmul(encoder, out, a, b, c, alpha_ptr, beta_ptr);
 }
 
 } // namespace mlx::core

mlx/backend/cuda/gemms/cublas_gemm.h

@@ -2,13 +2,14 @@
 #pragma once
 
 #include "mlx/array.h"
+#include "mlx/backend/cuda/cublas_utils.h"
 #include "mlx/backend/cuda/device.h"
 
 #include <cublasLt.h>
 
 namespace mlx::core {
 
-class CublasGemm {
+class CublasGemm : public CublasMatmulBase {
  public:
   CublasGemm(
       cu::Device& device,
@@ -42,8 +43,6 @@ class CublasGemm {
       int64_t b_batch_stride,
       int64_t c_batch_stride);
 
-  ~CublasGemm();
-
   // The output's descriptor is inferred from inputs by default, use this method
   // for unusual output.
   void set_out(
@@ -55,8 +54,6 @@ class CublasGemm {
       int32_t batch_count,
       int64_t batch_stride);
 
-  void set_bias(cu::CommandEncoder& encoder, const array& bias);
-
   void run(
       cu::CommandEncoder& encoder,
       array& out,
@@ -112,18 +109,6 @@ class CublasGemm {
       const void* c,
       float alpha = 1,
       float beta = 0);
-
-  uint64_t M_;
-  uint64_t N_;
-  cudaDataType_t scale_type_;
-  cublasLtMatmulPreference_t pref_{nullptr};
-  cublasLtHandle_t handle_{nullptr};
-  cublasLtMatmulDesc_t matmul_desc_{nullptr};
-  cublasLtMatrixLayout_t a_desc_{nullptr};
-  cublasLtMatrixLayout_t b_desc_{nullptr};
-  cublasLtMatrixLayout_t c_desc_{nullptr};
-  cublasLtMatrixLayout_t out_desc_{nullptr};
-  cublasLtMatmulHeuristicResult_t heuristic_;
 };
 
 } // namespace mlx::core

mlx/backend/cuda/gemms/grouped_gemm.h

@@ -0,0 +1,25 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+namespace mlx::core {
+
+namespace cu {
+class CommandEncoder;
+}
+
+class array;
+
+void cutlass_grouped_gemm_unaligned(
+    bool a_transposed,
+    int lda,
+    bool b_transposed,
+    int ldb,
+    int group_count,
+    const array& a,
+    const array& b,
+    const array& indices,
+    array& out,
+    cu::CommandEncoder& encoder);
+
+} // namespace mlx::core

mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu

@@ -0,0 +1,288 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/gemms/grouped_gemm.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/dtype_utils.h"
+
+#include <cooperative_groups.h>
+#include <cutlass/gemm/device/default_gemm_configuration.h>
+#include <cutlass/gemm/device/gemm_grouped.h>
+#include <cutlass/gemm/kernel/default_gemm_grouped.h>
+#include <fmt/format.h>
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+using ProblemSize = cutlass::gemm::GemmCoord;
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <int N_READS>
+__global__ void prepare_grouped_mm_data(
+    const uint32_t* indices,
+    size_t size,
+    int group_count,
+    int K,
+    int N,
+    int lda,
+    int ldb,
+    int item_size,
+    int8_t* a_start,
+    int8_t* b_start,
+    int8_t* out_start,
+    int a_batch_stride,
+    int b_batch_stride,
+    int out_batch_stride,
+    ProblemSize* problem_sizes,
+    int64_t* a_lds,
+    int64_t* b_lds,
+    int64_t* out_lds,
+    void** a_ptrs,
+    void** b_ptrs,
+    void** out_ptrs) {
+  auto block = cg::this_thread_block();
+
+  // cumsum(histogram(indices)) - offset for each group.
+  extern __shared__ uint32_t cum_histo[];
+
+  int group = block.thread_rank();
+  if (group < group_count) {
+    cum_histo[group] = 0;
+  }
+
+  block.sync();
+
+  // Since |indices| is sorted, the position where element changes would be its
+  // cumulative histogram.
+  size_t elems_per_block = block.num_threads() * N_READS;
+  for (int r = 0; r < cuda::ceil_div(size, elems_per_block); ++r) {
+    // TODO: Use vectorized read.
+    for (int i = 0; i < N_READS; ++i) {
+      size_t pos = r * elems_per_block + group * N_READS + i;
+      if (pos >= size) {
+        break;
+      }
+      auto elem = indices[pos];
+      auto next = pos < size - 1 ? indices[pos + 1] : group_count;
+      while (elem < next) {
+        cum_histo[elem] = pos + 1;
+        elem++;
+      }
+    }
+  }
+
+  block.sync();
+
+  if (group < group_count) {
+    // Fill shapes.
+    int delta =
+        group == 0 ? cum_histo[0] : cum_histo[group] - cum_histo[group - 1];
+    problem_sizes[group] = {delta, N, K};
+    a_lds[group] = lda;
+    b_lds[group] = ldb;
+    out_lds[group] = N;
+    // Fill pointers.
+    auto offset = group == 0 ? 0 : cum_histo[group - 1];
+    a_ptrs[group] = a_start + offset * item_size * a_batch_stride;
+    b_ptrs[group] = b_start + group * item_size * b_batch_stride;
+    out_ptrs[group] = out_start + offset * item_size * out_batch_stride;
+  }
+}
+
+} // namespace cu
+
+namespace {
+
+template <typename T, int kAlignment, typename Arch, typename OpClass>
+void grouped_gemm_v2(
+    bool a_transposed,
+    bool b_transposed,
+    int group_count,
+    ProblemSize* problem_sizes,
+    int64_t* a_lds,
+    int64_t* b_lds,
+    int64_t* out_lds,
+    void* a_ptrs,
+    void* b_ptrs,
+    void* out_ptrs,
+    cu::CommandEncoder& encoder) {
+  using ElementAccumulator = float;
+  using GemmConfiguration = typename cutlass::gemm::device::
+      DefaultGemmConfiguration<OpClass, Arch, T, T, T, ElementAccumulator>;
+  using EpilogueOutputOp = typename GemmConfiguration::EpilogueOutputOp;
+
+  dispatch_bool(a_transposed, [&](auto a_transposed_tag) {
+    dispatch_bool(b_transposed, [&](auto b_transposed_tag) {
+      using LayoutA = std::conditional_t<
+          a_transposed_tag,
+          cutlass::layout::ColumnMajor,
+          cutlass::layout::RowMajor>;
+      using LayoutB = std::conditional_t<
+          b_transposed_tag,
+          cutlass::layout::ColumnMajor,
+          cutlass::layout::RowMajor>;
+      using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmGrouped<
+          T,
+          LayoutA,
+          cutlass::ComplexTransform::kNone,
+          kAlignment,
+          T,
+          LayoutB,
+          cutlass::ComplexTransform::kNone,
+          kAlignment,
+          T,
+          cutlass::layout::RowMajor,
+          ElementAccumulator,
+          OpClass,
+          Arch,
+          typename GemmConfiguration::ThreadblockShape,
+          typename GemmConfiguration::WarpShape,
+          typename GemmConfiguration::InstructionShape,
+          EpilogueOutputOp,
+          cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,
+          GemmConfiguration::kStages>::GemmKernel;
+      using GemmGrouped =
+          typename cutlass::gemm::device::GemmGrouped<GemmKernel>;
+
+      typename EpilogueOutputOp::Params epilogue_op(
+          /* alpha */ 1, /* beta */ 0);
+      typename GemmGrouped::Arguments args(
+          problem_sizes,
+          group_count,
+          GemmGrouped::sufficient(),
+          epilogue_op,
+          reinterpret_cast<T**>(a_ptrs),
+          reinterpret_cast<T**>(b_ptrs),
+          reinterpret_cast<T**>(out_ptrs),
+          reinterpret_cast<T**>(out_ptrs),
+          a_lds,
+          b_lds,
+          out_lds,
+          out_lds);
+
+      GemmGrouped gemm;
+      cutlass::Status status = gemm.initialize(args, nullptr, encoder.stream());
+      if (status != cutlass::Status::kSuccess) {
+        throw std::runtime_error(fmt::format(
+            "Failed to initialize GemmGrouped: {}",
+            cutlass::cutlassGetStatusString(status)));
+      }
+
+      auto capture = encoder.capture_context();
+      status = gemm.run(encoder.stream());
+      if (status != cutlass::Status::kSuccess) {
+        throw std::runtime_error(fmt::format(
+            "Failed to run GemmGrouped: {}",
+            cutlass::cutlassGetStatusString(status)));
+      }
+    });
+  });
+}
+
+} // namespace
+
+void cutlass_grouped_gemm_unaligned(
+    bool a_transposed,
+    int lda,
+    bool b_transposed,
+    int ldb,
+    int group_count,
+    const array& a,
+    const array& b,
+    const array& indices,
+    array& out,
+    cu::CommandEncoder& encoder) {
+  // Prepare device pointers for matmul.
+  int problem_sizes_nbytes =
+      group_count * cuda::ceil_div(sizeof(ProblemSize), 8) * 8;
+  int nbytes = problem_sizes_nbytes +
+      group_count * (3 * sizeof(void*) + 3 * sizeof(int64_t));
+  nbytes = cuda::ceil_div(nbytes, 256) * 256;
+  array gemm_args(cu::malloc_async(nbytes, encoder), {nbytes}, int8);
+  encoder.add_temporary(gemm_args);
+
+  ProblemSize* problem_sizes = gpu_ptr<ProblemSize>(gemm_args);
+  int64_t* a_lds = gpu_ptr<int64_t>(gemm_args) + problem_sizes_nbytes / 8;
+  int64_t* b_lds = a_lds + group_count;
+  int64_t* out_lds = b_lds + group_count;
+  void** a_ptrs = reinterpret_cast<void**>(out_lds + group_count);
+  void** b_ptrs = a_ptrs + group_count;
+  void** out_ptrs = b_ptrs + group_count;
+
+  // Fill the pointers by computing offsets from indices.
+  constexpr int N_READS = 4;
+  size_t n_threads = cuda::ceil_div(indices.size(), N_READS);
+  n_threads = group_count < n_threads ? n_threads : group_count;
+  dim3 block_dims(std::min(n_threads, 1024ul));
+  dim3 num_blocks(1);
+
+  encoder.set_input_array(indices);
+  encoder.set_output_array(gemm_args);
+  encoder.add_kernel_node(
+      cu::prepare_grouped_mm_data<N_READS>,
+      num_blocks,
+      block_dims,
+      group_count * sizeof(uint32_t), // sizeof(cum_histo)
+      gpu_ptr<uint32_t>(indices),
+      indices.size(),
+      group_count,
+      a.shape(-1), // K
+      b.shape(-1), // N,
+      lda,
+      ldb,
+      out.itemsize(),
+      gpu_ptr<int8_t>(a),
+      gpu_ptr<int8_t>(b),
+      gpu_ptr<int8_t>(out),
+      a.shape(-2) * a.shape(-1), // a_batch_stride
+      b.shape(-2) * b.shape(-1), // b_batch_stride
+      out.shape(-2) * out.shape(-1), // out_batch_stride
+      problem_sizes,
+      a_lds,
+      b_lds,
+      out_lds,
+      a_ptrs,
+      b_ptrs,
+      out_ptrs);
+
+  // Invoke grouped GEMM.
+  constexpr int kAlignment = 1;
+  using Arch = cutlass::arch::Sm75;
+  using OpClass = cutlass::arch::OpClassSimt;
+  auto* fun = grouped_gemm_v2<float, kAlignment, Arch, OpClass>;
+  switch (a.dtype()) {
+    case float32:
+      break;
+    case float16:
+      fun = grouped_gemm_v2<cutlass::half_t, kAlignment, Arch, OpClass>;
+      break;
+    case bfloat16:
+      fun = grouped_gemm_v2<cutlass::bfloat16_t, kAlignment, Arch, OpClass>;
+      break;
+    default:
+      throw std::runtime_error(fmt::format(
+          "Unsupported dtype in cutlass_grouped_gemm_sm75: {}.",
+          dtype_to_string(a.dtype())));
+  }
+
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(gemm_args);
+  encoder.set_output_array(out);
+  fun(a_transposed,
+      b_transposed,
+      group_count,
+      problem_sizes,
+      a_lds,
+      b_lds,
+      out_lds,
+      a_ptrs,
+      b_ptrs,
+      out_ptrs,
+      encoder);
+}
+
+} // namespace mlx::core

mlx/backend/cuda/jit_module.cpp

@@ -27,55 +27,47 @@ void check_nvrtc_error(const char* name, nvrtcResult err) {
   }
 }
 
-// Return the location of the CUDA toolkit.
-const std::string& cuda_home() {
-  static std::string home = []() -> std::string {
-    const char* home = std::getenv("CUDA_HOME");
-    if (home) {
-      return home;
-    }
-    home = std::getenv("CUDA_PATH");
-    if (home) {
-      return home;
-    }
-#if defined(__linux__)
-    home = "/usr/local/cuda";
-    if (std::filesystem::exists(home)) {
-      return home;
-    }
-#endif
-    throw std::runtime_error(
-        "Environment variable CUDA_HOME or CUDA_PATH is not set.");
-  }();
-  return home;
-}
-
-// Return the location of CCCL headers shipped with the distribution.
-const std::string& cccl_dir() {
-  static std::string dir = []() {
-    std::filesystem::path path;
+// Return the --include-path args used for invoking NVRTC.
+const std::vector<std::string>& include_path_args() {
+  static std::vector<std::string> cached_args = []() {
+    std::vector<std::string> args;
+    // Add path to bundled CCCL headers.
+    auto root_dir = current_binary_dir().parent_path();
+    auto path = root_dir / "include" / "cccl";
 #if defined(MLX_CCCL_DIR)
-    // First search the install dir if defined.
-    path = MLX_CCCL_DIR;
-    if (std::filesystem::exists(path)) {
-      return path.string();
+    if (!std::filesystem::exists(path)) {
+      path = MLX_CCCL_DIR;
     }
 #endif
-    // Then search dynamically from the dir of libmlx.so file.
-    path = current_binary_dir().parent_path() / "include" / "cccl";
     if (std::filesystem::exists(path)) {
-      return path.string();
+      args.push_back(fmt::format("--include-path={}", path.string()));
     }
-    // Finally check the environment variable.
-    if (const char* env = std::getenv("MLX_CCCL_DIR"); env) {
-      path = env;
-      if (!path.empty() && std::filesystem::exists(path)) {
-        return path.string();
+    // Add path to CUDA runtime headers, try local-installed python package
+    // first and then system-installed headers.
+    path = root_dir.parent_path() / "nvidia" / "cuda_runtime" / "include";
+    if (std::filesystem::exists(path)) {
+      args.push_back(fmt::format("--include-path={}", path.string()));
+    } else {
+      const char* home = std::getenv("CUDA_HOME");
+      if (!home) {
+        home = std::getenv("CUDA_PATH");
+      }
+#if defined(__linux__)
+      if (!home) {
+        home = "/usr/local/cuda";
+      }
+#endif
+      if (home && std::filesystem::exists(home)) {
+        args.push_back(fmt::format("--include-path={}/include", home));
+      } else {
+        throw std::runtime_error(
+            "Can not find locations of CUDA headers, please set environment "
+            "variable CUDA_HOME or CUDA_PATH.");
       }
     }
-    return std::string();
+    return args;
   }();
-  return dir;
+  return cached_args;
 }
 
 // Get the cache directory for storing compiled results.
@@ -288,14 +280,9 @@ void compile(
       device.compute_capability_minor(),
       arch_tag);
   args.push_back(compute.c_str());
-  std::string cccl_include = cccl_dir();
-  if (!cccl_include.empty()) {
-    cccl_include = fmt::format("--include-path={}", cccl_include);
-    args.push_back(cccl_include.c_str());
+  for (const auto& include : include_path_args()) {
+    args.push_back(include.c_str());
   }
-  std::string cuda_include =
-      fmt::format("--include-path={}/include", cuda_home());
-  args.push_back(cuda_include.c_str());
   nvrtcResult compile_result =
       nvrtcCompileProgram(prog, args.size(), args.data());
   if (compile_result != NVRTC_SUCCESS) {

mlx/backend/cuda/matmul.cpp

@@ -4,6 +4,7 @@
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/gemms/cublas_gemm.h"
 #include "mlx/backend/cuda/gemms/gemv.h"
+#include "mlx/backend/cuda/gemms/grouped_gemm.h"
 #include "mlx/backend/gpu/copy.h"
 #include "mlx/primitives.h"
 
@@ -29,6 +30,38 @@ check_transpose(cu::CommandEncoder& enc, const Stream& s, const array& arr) {
   }
 }
 
+std::tuple<bool, int64_t, array>
+ensure_batch_contiguous(const array& x, cu::CommandEncoder& encoder, Stream s) {
+  if (x.flags().row_contiguous) {
+    return std::make_tuple(false, x.strides(-2), x);
+  }
+
+  bool rc = true;
+  for (int i = 0; i < x.ndim() - 3; i++) {
+    rc &= (x.strides(i + 1) * x.shape(i)) == x.strides(i);
+  }
+  if (rc) {
+    return check_transpose(encoder, s, x);
+  }
+
+  array x_copy = contiguous_copy_gpu(x, s);
+  encoder.add_temporary(x_copy);
+  return std::make_tuple(false, x_copy.strides(-2), x_copy);
+}
+
+array ensure_row_contiguous(
+    const array& x,
+    cu::CommandEncoder& encoder,
+    Stream s) {
+  if (!x.flags().row_contiguous) {
+    array x_copy = contiguous_copy_gpu(x, s);
+    encoder.add_temporary(x_copy);
+    return x_copy;
+  } else {
+    return x;
+  }
+}
+
 void gemm_and_bias(
     cu::CommandEncoder& encoder,
     int M,
@@ -103,6 +136,40 @@ void gemm_and_bias(
       encoder, out, a, b, batch_shape, a_batch_strides, b_batch_strides, alpha);
 }
 
+void gather_mm_rhs(
+    const array& a_,
+    const array& b_,
+    const array& indices_,
+    array& out,
+    cu::CommandEncoder& encoder,
+    Stream s) {
+  if (a_.size() / a_.shape(-2) / a_.shape(-1) != indices_.size()) {
+    throw std::runtime_error("[gather_mm] Broadcasting lhs is not supported.");
+  }
+
+  int group_count = b_.size() / b_.shape(-1) / b_.shape(-2);
+  if (group_count > 1024) {
+    throw std::runtime_error(
+        "[gather_mm] Group count can not be larger than 1024.");
+  }
+
+  auto [a_transposed, lda, a] = ensure_batch_contiguous(a_, encoder, s);
+  auto [b_transposed, ldb, b] = ensure_batch_contiguous(b_, encoder, s);
+  auto indices = ensure_row_contiguous(indices_, encoder, s);
+
+  cutlass_grouped_gemm_unaligned(
+      a_transposed,
+      lda,
+      b_transposed,
+      ldb,
+      group_count,
+      a,
+      b,
+      indices,
+      out,
+      encoder);
+}
+
 } // namespace
 
 void Matmul::eval_gpu(const std::vector<array>& inputs, array& out) {
@@ -254,4 +321,40 @@ void AddMM::eval_gpu(const std::vector<array>& inputs, array& out) {
       beta_);
 }
 
+void GatherMM::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("GatherMM::eval_gpu");
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  assert(inputs.size() == 4);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  auto& lhs_indices = inputs[2];
+  auto& rhs_indices = inputs[3];
+
+  // Return 0s if either input is empty.
+  if (a.size() == 0 || b.size() == 0) {
+    array zero(0, a.dtype());
+    encoder.add_temporary(zero);
+    fill_gpu(zero, out, s);
+    return;
+  }
+
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+
+  // Extract shapes from inputs.
+  int M = a.shape(-2);
+  int N = b.shape(-1);
+  int K = a.shape(-1);
+
+  // We are walking a in order and b is also in order so we can batch up the
+  // matmuls and reuse reading a and b.
+  if (M == 1 && right_sorted_ == true) {
+    gather_mm_rhs(a, b, rhs_indices, out, encoder, s);
+    return;
+  }
+
+  throw std::runtime_error("NYI");
+}
+
 } // namespace mlx::core

mlx/backend/cuda/primitives.cpp

@@ -1,6 +1,7 @@
 // Copyright © 2025 Apple Inc.
 
 #include "mlx/distributed/primitives.h"
+#include <cuda_runtime.h>
 #include "mlx/fast_primitives.h"
 #include "mlx/primitives.h"
 
@@ -23,9 +24,15 @@ namespace mlx::core {
     throw std::runtime_error(#func " has no CUDA implementation.");   \
   }
 
+#if CUDART_VERSION < 12080
+void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
+  throw std::runtime_error(
+      "[QQMatmul::eval_gpu] QQMM is only supported with CUDA 12.8 or higher.");
+}
+#endif
+
 NO_GPU(BlockMaskedMM)
 NO_GPU(FFT)
-NO_GPU(GatherMM)
 NO_GPU(GatherQMM)
 NO_GPU(Hadamard)
 NO_GPU_MULTI(LUF)

mlx/backend/cuda/quantized/cublas_qqmm.cpp

@@ -0,0 +1,206 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/quantized/cublas_qqmm.h"
+
+#include <fmt/format.h>
+#include "mlx/backend/cuda/cublas_utils.h"
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/dtype_utils.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+namespace {
+
+// Currently cublas supports only mxfp8 and nvfp4
+// quantization modes for block scaled quantization
+cudaDataType_t qmode_to_cublas_scale_dtype(std::string mode) {
+  if (mode == "mxfp8") {
+    return CUDA_R_8F_UE8M0;
+  } else if (mode == "nvfp4") {
+    return CUDA_R_8F_UE4M3;
+  } else {
+    throw std::runtime_error(
+        fmt::format("Unsupported quantization mode in CublasQQMM: {}.", mode));
+  }
+}
+
+cudaDataType_t qmode_to_cublas_dtype(std::string mode) {
+  if (mode == "mxfp8") {
+    return CUDA_R_8F_E4M3;
+  } else if (mode == "nvfp4") {
+    return CUDA_R_4F_E2M1;
+  } else {
+    throw std::runtime_error(
+        fmt::format("Unsupported quantization mode in CublasQQMM: {}.", mode));
+  }
+}
+
+cublasLtMatmulMatrixScale_t qmode_to_cublas_scale_mode(std::string mode) {
+  if (mode == "mxfp8") {
+    return CUBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0;
+  } else if (mode == "nvfp4") {
+    return CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3;
+  } else {
+    throw std::runtime_error(
+        fmt::format("Unsupported quantization mode in CublasQQMM: {}.", mode));
+  }
+}
+
+} // namespace
+
+CublasQQMM::CublasQQMM(
+    cu::Device& device,
+    bool a_transposed,
+    uint64_t a_rows,
+    uint64_t a_cols,
+    int64_t lda,
+    bool b_transposed,
+    uint64_t b_rows,
+    uint64_t b_cols,
+    int64_t ldb,
+    int32_t batch_count,
+    int64_t a_batch_stride,
+    int64_t b_batch_stride,
+    Dtype out_dtype,
+    std::string qmode) {
+  // The compute type must be CUBLAS_COMPUTE_32F.
+  // The scale type must be CUDA_R_32F.
+  cudaDataType_t scale_type = CUDA_R_32F;
+  cublasComputeType_t gemm_compute_type = CUBLAS_COMPUTE_32F;
+  cudaDataType_t output_type =
+      cublas_utils::dtype_to_cublas_type(out_dtype, "CublasQQMM");
+  cudaDataType_t data_type = qmode_to_cublas_dtype(qmode);
+  quantization_mode_ = std::string(qmode);
+
+  init_base(
+      device,
+      scale_type,
+      gemm_compute_type,
+      data_type,
+      output_type,
+      a_transposed,
+      a_rows,
+      a_cols,
+      lda,
+      b_transposed,
+      b_rows,
+      b_cols,
+      ldb,
+      batch_count,
+      a_batch_stride,
+      b_batch_stride);
+
+  a_scale_mode_ = qmode_to_cublas_scale_mode(qmode);
+  b_scale_mode_ = qmode_to_cublas_scale_mode(qmode);
+
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_B_SCALE_MODE,
+      &a_scale_mode_,
+      sizeof(a_scale_mode_)));
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_A_SCALE_MODE,
+      &b_scale_mode_,
+      sizeof(b_scale_mode_)));
+}
+
+CublasQQMM::CublasQQMM(
+    cu::Device& device,
+    bool a_transposed,
+    uint64_t a_rows,
+    uint64_t a_cols,
+    int64_t lda,
+    bool b_transposed,
+    uint64_t b_rows,
+    uint64_t b_cols,
+    int64_t ldb,
+    int64_t ldc,
+    int32_t batch_count,
+    int64_t a_batch_stride,
+    int64_t b_batch_stride,
+    int64_t c_batch_stride,
+    Dtype out_dtype,
+    std::string qmode)
+    : CublasQQMM(
+          device,
+          a_transposed,
+          a_rows,
+          a_cols,
+          lda,
+          b_transposed,
+          b_rows,
+          b_cols,
+          ldb,
+          batch_count,
+          a_batch_stride,
+          b_batch_stride,
+          out_dtype,
+          qmode) {
+  auto type = cublas_utils::dtype_to_cublas_type(
+      out_dtype, "CublasQQMM"); // must match the output type
+  c_desc_ = cublas_utils::create_matrix_layout(
+      type,
+      b_transposed ? b_rows : b_cols,
+      a_transposed ? a_cols : a_rows,
+      false,
+      ldc,
+      batch_count,
+      c_batch_stride);
+}
+
+void CublasQQMM::run(
+    cu::CommandEncoder& encoder,
+    array& out,
+    const array& a,
+    const array& b,
+    const array& a_scale,
+    const array& b_scale,
+    float alpha) {
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(a_scale);
+  encoder.set_input_array(b_scale);
+  encoder.set_output_array(out);
+
+  execute(
+      encoder,
+      gpu_ptr<void>(out),
+      gpu_ptr<void>(a),
+      gpu_ptr<void>(b),
+      gpu_ptr<void>(a_scale),
+      gpu_ptr<void>(b_scale),
+      nullptr,
+      alpha);
+}
+
+void CublasQQMM::execute(
+    cu::CommandEncoder& encoder,
+    void* out,
+    const void* a,
+    const void* b,
+    const void* a_scale,
+    const void* b_scale,
+    const void* c,
+    float alpha /* = 1 */,
+    float beta /* = 0 */) {
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,
+      &b_scale,
+      sizeof(b_scale)));
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
+      &a_scale,
+      sizeof(a_scale)));
+
+  const void* alpha_ptr = &alpha;
+  const void* beta_ptr = &beta;
+
+  execute_matmul(encoder, out, a, b, c, alpha_ptr, beta_ptr);
+}
+
+} // namespace mlx::core

mlx/backend/cuda/quantized/cublas_qqmm.h

@@ -0,0 +1,88 @@
+// Copyright © 2025 Apple Inc.
+#pragma once
+
+#include "mlx/array.h"
+#include "mlx/backend/cuda/cublas_utils.h"
+#include "mlx/backend/cuda/device.h"
+
+#include <cublasLt.h>
+
+namespace mlx::core {
+
+class CublasQQMM : public CublasMatmulBase {
+ public:
+  CublasQQMM(
+      cu::Device& device,
+      bool a_transposed,
+      uint64_t a_rows,
+      uint64_t a_cols,
+      int64_t lda,
+      bool b_transposed,
+      uint64_t b_rows,
+      uint64_t b_cols,
+      int64_t ldb,
+      int32_t batch_count,
+      int64_t a_batch_stride,
+      int64_t b_batch_stride,
+      Dtype out_dtype,
+      std::string quantization_mode);
+
+  CublasQQMM(
+      cu::Device& device,
+      bool a_transposed,
+      uint64_t a_rows,
+      uint64_t a_cols,
+      int64_t lda,
+      bool b_transposed,
+      uint64_t b_rows,
+      uint64_t b_cols,
+      int64_t ldb,
+      int64_t ldc,
+      int32_t batch_count,
+      int64_t a_batch_stride,
+      int64_t b_batch_stride,
+      int64_t c_batch_stride,
+      Dtype out_dtype,
+      std::string quantization_mode);
+
+  void run(
+      cu::CommandEncoder& encoder,
+      array& out,
+      const array& a,
+      const array& b,
+      const array& a_scale,
+      const array& b_scale,
+      float alpha = 1.0f);
+
+ private:
+  void run_batched(
+      cu::CommandEncoder& encoder,
+      array& out,
+      const array& a,
+      const array& b,
+      const array& a_scale,
+      const array& b_scale,
+      const Shape& batch_shape,
+      const Strides& a_batch_strides,
+      const Strides& b_batch_strides,
+      float alpha);
+
+  void execute(
+      cu::CommandEncoder& encoder,
+      void* out,
+      const void* a,
+      const void* b,
+      const void* a_scale,
+      const void* b_scale,
+      const void* c,
+      float alpha = 1,
+      float beta = 0);
+
+  std::string quantization_mode_;
+  cublasLtMatmulMatrixScale_t a_scale_mode_;
+  cublasLtMatmulMatrixScale_t b_scale_mode_;
+  cublasLtMatmulMatrixScale_t c_scale_mode_;
+  cublasLtMatmulMatrixScale_t out_scale_mode_;
+};
+
+} // namespace mlx::core

mlx/backend/cuda/quantized/fp_quantize.cu

@@ -2,7 +2,11 @@
 
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/cuda/quantized/mxfp8_quantize.cuh"
+#include "mlx/backend/cuda/quantized/nvfp4_quantize.cuh"
 #include "mlx/backend/cuda/quantized/quantized.h"
+#include "mlx/backend/cuda/quantized/quantized_utils.cuh"
+#include "mlx/backend/cuda/vector_types.cuh"
 #include "mlx/dtype_utils.h"
 
 #include <cooperative_groups.h>
@@ -13,17 +17,6 @@
 namespace mlx::core {
 namespace cu {
 
-template <int bits>
-struct Quantize {
-  __device__ uint8_t operator()(float x) {
-    if constexpr (bits == 8) {
-      return __nv_fp8_e4m3(x).__x;
-    } else {
-      return __nv_fp4_e2m1(x).__x;
-    }
-  }
-};
-
 template <int bits>
 struct Dequantize {
   __device__ float operator()(uint8_t x) {
@@ -37,29 +30,40 @@ struct Dequantize {
 
 namespace cg = cooperative_groups;
 
-template <typename T, int group_size, int bits, bool use_mx_scale>
-__global__ void
-fp_quantize(const T* w, uint8_t* out, uint8_t* scales, size_t size) {
+template <typename T, int group_size, int bits, bool use_mx_scale, bool USE_SR>
+__global__ void fp_quantize(T* w, uint8_t* out, uint8_t* scales, size_t size) {
+  using Tx2 = Vector2_t<T>;
+  using Tx4 = Vector4_t<T>;
+  uint32_t rbits = 0; // reserved bits for future use
   auto block_size = cg::this_thread_block().dim_threads();
   auto block_idx = cg::this_thread_block().group_index();
   auto idx_in_block = cg::this_thread_block().thread_index();
-
   auto tidx = block_idx.x * block_size.x + idx_in_block.x;
   auto tidy = block_idx.y * block_size.y + idx_in_block.y;
+  auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x;
 
-  auto grid_dim_x =
-      cg::this_grid().dim_blocks().x * cg::this_grid().block_index().x;
-  size_t index = tidx + grid_dim_x * size_t(tidy);
-  if (index >= size) {
+  size_t thread_idx = tidx + grid_dim_x * size_t(tidy);
+  size_t base_idx = thread_idx * group_size;
+
+  if (base_idx >= size) {
     return;
   }
 
-  float w_thread = w[index];
+  auto w_tile = load_vector<group_size, T>(w, thread_idx);
+  float scale = 0.0f;
 
-  cg::greater<float> max_op;
-  auto warp = cg::tiled_partition<group_size>(cg::this_thread_block());
+  Tx2 amax_2x = Tx2{0.0f, 0.0f};
+
+#pragma unroll
+  for (int i = 0; i < group_size; i += 2) {
+    auto pair = Tx2{w_tile[i], w_tile[i + 1]};
+    abs_max_x2<Tx2>(amax_2x, amax_2x, pair);
+  }
+
+  scale = static_cast<float>(
+      max(fabsf(static_cast<float>(amax_2x.x)),
+          fabsf(static_cast<float>(amax_2x.y))));
 
-  float scale = cg::reduce(warp, abs(w_thread), max_op);
   scale /= bits == 4 ? 6.0f : 448.0f;
   // Convert to mx scale or nv scale
   using ScaleType =
@@ -68,21 +72,24 @@ fp_quantize(const T* w, uint8_t* out, uint8_t* scales, size_t size) {
   uint8_t q_scale = s.__x;
   scale = float(s);
 
-  // Write out the scales
-  size_t gindex = index / group_size;
-  if (index % group_size == 0) {
-    scales[gindex] = q_scale;
-  }
+  scales[thread_idx] = q_scale;
+  constexpr int elem_per_byte = bits == 8 ? 1 : 2;
+  AlignedVector<uint8_t, group_size / elem_per_byte> quantized;
 
-  uint8_t output = Quantize<bits>{}(scale == 0 ? 0.0f : w_thread / scale);
-  if (bits == 4) {
-    uint8_t sval = warp.shfl_down(output, 1);
-    output |= sval << bits;
-  }
-  constexpr int pack_factor = bits == 8 ? 1 : 2;
-  if (index % pack_factor == 0) {
-    out[index / pack_factor] = output;
+#pragma unroll
+  for (int i = 0; i < group_size / 4; i++) {
+    Tx4 w_Tx4 = *reinterpret_cast<Tx4*>(&w_tile[i * 4]);
+    if constexpr (bits == 8) {
+      uint32_t quantized_val =
+          scale_cvt_Tx4_to_fp8x4<T, USE_SR>(w_Tx4, 1.0f / scale, rbits);
+      *reinterpret_cast<uint32_t*>(&quantized[i * 4]) = quantized_val;
+    } else {
+      uint16_t quantized_val =
+          scale_cvt_Tx4_to_fp4x4<T, USE_SR>(w_Tx4, 1.0f / scale, rbits);
+      *reinterpret_cast<uint16_t*>(&quantized[i * 2]) = quantized_val;
+    }
   }
+  store_vector<group_size / elem_per_byte>(out, thread_idx, quantized);
 }
 
 template <typename T, int group_size, int bits, bool use_mx_scale>
@@ -95,8 +102,7 @@ fp_dequantize(const uint8_t* w, const uint8_t* scales, T* out, size_t size) {
   auto tidx = block_idx.x * block_size.x + idx_in_block.x;
   auto tidy = block_idx.y * block_size.y + idx_in_block.y;
 
-  auto grid_dim_x =
-      cg::this_grid().dim_blocks().x * cg::this_grid().block_index().x;
+  auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x;
 
   constexpr int pack_factor = bits == 8 ? 1 : 2;
   size_t offset = tidx + grid_dim_x * size_t(tidy);
@@ -142,15 +148,16 @@ void fp_quantize(
   dispatch_float_types(w.dtype(), "fp_quantize", [&](auto type_tag) {
     using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
     if constexpr (!std::is_same_v<T, double>) {
-      auto kernel = cu::fp_quantize<T, 32, 4, true>;
+      auto kernel = cu::fp_quantize<T, 32, 4, true, false>;
       if (bits == 8) {
-        kernel = cu::fp_quantize<T, 32, 8, true>;
+        kernel = cu::fp_quantize<T, 32, 8, true, false>;
       } else if (group_size == 16) {
-        kernel = cu::fp_quantize<T, 16, 4, false>;
+        kernel = cu::fp_quantize<T, 16, 4, false, false>;
       }
       bool large = w.size() > UINT_MAX;
       auto [num_blocks, block_dims] =
-          get_launch_args(w.size(), w.shape(), w.strides(), large);
+          get_launch_args(w.size(), w.shape(), w.strides(), large, group_size);
+
       enc.add_kernel_node(
           kernel,
           num_blocks,

mlx/backend/cuda/quantized/mxfp8_quantize.cuh

@@ -0,0 +1,32 @@
+#pragma once
+
+#include <cuda.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#include "mlx/backend/cuda/vector_types.cuh"
+
+namespace mlx::core::cu {
+
+// TODO implement fast path
+template <typename T>
+__device__ __forceinline__ uint32_t
+scale_cvt_Tx4_to_fp8x4_fallback(const Vector4_t<T> input, const float scale) {
+  uint32_t out_fp8x4 = 0;
+  float4 scaled;
+  scaled.x = static_cast<float>(input.x) * scale;
+  scaled.y = static_cast<float>(input.y) * scale;
+  scaled.z = static_cast<float>(input.z) * scale;
+  scaled.w = static_cast<float>(input.w) * scale;
+  out_fp8x4 = __nv_fp8x4_e4m3(scaled).__x;
+  return out_fp8x4;
+}
+
+// Place holder for future fast path implementation
+template <typename T, bool USE_SR>
+__device__ __forceinline__ uint32_t scale_cvt_Tx4_to_fp8x4(
+    const Vector4_t<T> input,
+    const float scale,
+    uint32_t rbits) {
+  return scale_cvt_Tx4_to_fp8x4_fallback(input, scale);
+}
+} // namespace mlx::core::cu

mlx/backend/cuda/quantized/nvfp4_quantize.cuh

@@ -0,0 +1,334 @@
+#pragma once
+
+#include <cuda.h>
+#include <cuda_fp4.h>
+#include <cuda_runtime.h>
+#include "mlx/backend/cuda/vector_types.cuh"
+
+namespace mlx::core::cu {
+
+using bf16x4 = Vector4_t<__nv_bfloat16>;
+using fp16x4 = Vector4_t<__half>;
+using f32x4 = Vector4_t<float>;
+
+template <typename T>
+__device__ __forceinline__ uint16_t
+scale_cvt_Tx4_to_fp4x4_fallback(const Vector4_t<T> input, const float scale) {
+  // Fallback implementation for architectures that do not support cvt
+  // instructions or for cuda versions with no fp4 support (< 12.8) -> scalar
+  uint16_t out_fp4x4 = 0;
+  fp32x4 scaled;
+  scaled.x = static_cast<float>(input.x) * scale;
+  scaled.y = static_cast<float>(input.y) * scale;
+  scaled.z = static_cast<float>(input.z) * scale;
+  scaled.w = static_cast<float>(input.w) * scale;
+  uint8_t q0 = __nv_fp4_e2m1(scaled.x).__x;
+  uint8_t q1 = __nv_fp4_e2m1(scaled.y).__x;
+  uint8_t q2 = __nv_fp4_e2m1(scaled.z).__x;
+  uint8_t q3 = __nv_fp4_e2m1(scaled.w).__x;
+  out_fp4x4 = (static_cast<uint16_t>(q3) << 12) |
+      (static_cast<uint16_t>(q2) << 8) | (static_cast<uint16_t>(q1) << 4) |
+      static_cast<uint16_t>(q0);
+  return out_fp4x4;
+}
+
+#if (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) && \
+    defined(__CUDA_ARCH_SPECIFIC__)
+
+__device__ __forceinline__ uint16_t
+scale_cvt_bf16x4_to_fp4x4_rn(const bf16x4 input_bf16x4, const float2 scale) {
+  uint16_t out_fp4x4 = 0;
+  asm volatile(
+      "{\n"
+      ".reg.b16 x0_bf16; \n\t" // first bf16
+      ".reg.b16 x1_bf16; \n\t" // second bf16
+      ".reg.b16 x2_bf16; \n\t" // third bf16
+      ".reg.b16 x3_bf16; \n\t" // fourth bf16
+      ".reg.b32 x0; \n\t" // to hold scaled first
+      ".reg.b32 x1; \n\t" // to hold scaled second
+      ".reg.b32 x2; \n\t" // to hold scaled third
+      ".reg.b32 x3; \n\t" // to hold scaled fourth
+      ".reg.b64 x01; \n\t" // to hold vector mul
+      ".reg.b64 x23; \n\t"
+      ".reg.b8 q0; \n\t" // output byte fp4x2 (first pair)
+      ".reg.b8 q1; \n\t" // output byte fp4x2 (second pair)
+      "mov.b64 {x0_bf16, x1_bf16, x2_bf16, x3_bf16} , %1; \n\t" // unpack bf16
+      "cvt.f32.bf16 x0, x0_bf16; \n\t" // convert to f32
+      "cvt.f32.bf16 x1, x1_bf16; \n\t"
+      "cvt.f32.bf16 x2, x2_bf16; \n\t"
+      "cvt.f32.bf16 x3, x3_bf16; \n\t"
+      "mov.b64 x01, {x0, x1}; \n\t"
+      "mul.f32x2 x01, x01, %2; \n\t" // scale first pair
+      "mov.b64 x23, {x2, x3}; \n\t"
+      "mul.f32x2 x23, x23, %2; \n\t" // scale second pair
+      "mov.b64 {x0, x1}, x01; \n\t"
+      "mov.b64 {x2, x3}, x23; \n\t"
+      "cvt.rn.satfinite.e2m1x2.f32 q0, x1, x0; \n\t" // convert to fp4x2 first
+                                                     // pair
+      "cvt.rn.satfinite.e2m1x2.f32 q1, x3, x2; \n\t" // convert to fp4x2 second
+                                                     // pair
+      "mov.b16 %0, {q0, q1}; \n\t" // pack to output
+      "}"
+      : "=h"(out_fp4x4)
+      : "l"(reinterpret_cast<const uint64_t&>(input_bf16x4)),
+        "l"(reinterpret_cast<const uint64_t&>(
+            scale))); // here cast is needed becuase an asm operand must have
+                      // scalar type
+  return out_fp4x4;
+}
+
+__device__ __forceinline__ uint16_t scale_cvt_bf16x4_to_fp4x4_rs(
+    const bf16x4 input_bf16x4,
+    const float2 scale,
+    uint32_t rbits) {
+  uint16_t out_fp4x4 = 0;
+  asm volatile(
+      "{\n"
+      ".reg.b16 x0_bf16; \n\t"
+      ".reg.b16 x1_bf16; \n\t"
+      ".reg.b16 x2_bf16; \n\t"
+      ".reg.b16 x3_bf16; \n\t"
+      ".reg.b32 x0; \n\t"
+      ".reg.b32 x1; \n\t"
+      ".reg.b32 x2; \n\t"
+      ".reg.b32 x3; \n\t"
+      ".reg.b64 x01; \n\t"
+      ".reg.b64 x23; \n\t"
+      ".reg.b16 q0; \n\t"
+      "mov.b64 {x0_bf16, x1_bf16, x2_bf16, x3_bf16} , %1; \n\t"
+      "cvt.f32.bf16 x0, x0_bf16; \n\t"
+      "cvt.f32.bf16 x1, x1_bf16; \n\t"
+      "cvt.f32.bf16 x2, x2_bf16; \n\t"
+      "cvt.f32.bf16 x3, x3_bf16; \n\t"
+      "mov.b64 x01, {x0, x1}; \n\t"
+      "mul.f32x2 x01, x01, %2; \n\t"
+      "mov.b64 x23, {x2, x3}; \n\t"
+      "mul.f32x2 x23, x23, %2; \n\t"
+      "mov.b64 {x0, x1}, x01; \n\t"
+      "mov.b64 {x2, x3}, x23; \n\t"
+      "cvt.rs.satfinite.e2m1x4.f32 q0, {x3, x2, x1, x0}, %3; \n\t"
+      "}"
+      : "=h"(out_fp4x4)
+      : "l"(reinterpret_cast<const uint64_t&>(input_bf16x4)),
+        "l"(reinterpret_cast<const uint64_t&>(scale)),
+        "r"(rbits));
+  return out_fp4x4;
+}
+
+__device__ __forceinline__ uint16_t scale_cvt_fp32x4_to_fp4x4_rn(
+    const float2 input_fp32x2_0,
+    const float2 input_fp32x2_1,
+    const float2 scale) {
+  uint16_t out_fp4x4 = 0;
+  asm volatile(
+      "{\n"
+      ".reg.b32 x0; \n\t"
+      ".reg.b32 x1; \n\t"
+      ".reg.b32 x2; \n\t"
+      ".reg.b32 x3; \n\t"
+      ".reg.b64 x01; \n\t"
+      ".reg.b64 x23; \n\t"
+      ".reg.b8 q0; \n\t"
+      ".reg.b8 q1; \n\t"
+      "mov.b64 x01, {%1, %2}; \n\t"
+      "mul.f32x2 x01, x01, %5; \n\t"
+      "mov.b64 x23, {%3, %4}; \n\t"
+      "mul.f32x2 x23, x23, %5; \n\t"
+      "mov.b64 {x0, x1}, x01; \n\t"
+      "mov.b64 {x2, x3}, x23; \n\t"
+      "cvt.rn.satfinite.e2m1x2.f32 q0, x1, x0; \n\t"
+      "cvt.rn.satfinite.e2m1x2.f32 q1, x3, x2; \n\t"
+      "mov.b16 %0, {q0, q1}; \n\t"
+      "}"
+      : "=h"(out_fp4x4)
+      : "f"(input_fp32x2_0.x),
+        "f"(input_fp32x2_0.y),
+        "f"(input_fp32x2_1.x),
+        "f"(input_fp32x2_1.y),
+        "l"(reinterpret_cast<const uint64_t&>(scale)));
+  return out_fp4x4;
+}
+
+__device__ __forceinline__ uint16_t scale_cvt_fp32x4_to_fp4x4_rs(
+    const float2 input_fp32x2_0,
+    const float2 input_fp32x2_1,
+    const float2 scale,
+    uint32_t rbits) {
+  uint16_t out_fp4x4 = 0;
+  asm volatile(
+      "{\n"
+      ".reg.b32 x0; \n\t"
+      ".reg.b32 x1; \n\t"
+      ".reg.b32 x2; \n\t"
+      ".reg.b32 x3; \n\t"
+      ".reg.b64 x01; \n\t"
+      ".reg.b64 x23; \n\t"
+      ".reg.b16 q0; \n\t"
+      "mov.b64 x01, {%1, %2}; \n\t"
+      "mul.f32x2 x01, x01, %5; \n\t"
+      "mov.b64 x23, {%3, %4}; \n\t"
+      "mul.f32x2 x23, x23, %5; \n\t"
+      "mov.b64 {x0, x1}, x01; \n\t"
+      "mov.b64 {x2, x3}, x23; \n\t"
+      "cvt.rs.satfinite.e2m1x4.f32 q0, {x3, x2, x1, x0}, %6; \n\t"
+      "}"
+      : "=h"(out_fp4x4)
+      : "f"(input_fp32x2_0.x),
+        "f"(input_fp32x2_0.y),
+        "f"(input_fp32x2_1.x),
+        "f"(input_fp32x2_1.y),
+        "l"(reinterpret_cast<const uint64_t&>(scale)),
+        "r"(rbits));
+  return out_fp4x4;
+}
+
+__device__ __forceinline__ uint16_t
+scale_cvt_fp16x4_to_fp4x4_rn(const fp16x4 input_fp16x4, const float2 scale) {
+  uint16_t out_fp4x4 = 0;
+  asm volatile(
+      "{\n"
+      ".reg.b16 x0_fp16; \n\t"
+      ".reg.b16 x1_fp16; \n\t"
+      ".reg.b16 x2_fp16; \n\t"
+      ".reg.b16 x3_fp16; \n\t"
+      ".reg.b32 x0; \n\t"
+      ".reg.b32 x1; \n\t"
+      ".reg.b32 x2; \n\t"
+      ".reg.b32 x3; \n\t"
+      ".reg.b64 x01; \n\t"
+      ".reg.b64 x23; \n\t"
+      ".reg.b8 q0; \n\t"
+      ".reg.b8 q1; \n\t"
+      "mov.b64 {x0_fp16, x1_fp16, x2_fp16, x3_fp16} , %1; \n\t"
+      "cvt.f32.f16 x0, x0_fp16; \n\t"
+      "cvt.f32.f16 x1, x1_fp16; \n\t"
+      "cvt.f32.f16 x2, x2_fp16; \n\t"
+      "cvt.f32.f16 x3, x3_fp16; \n\t"
+      "mov.b64 x01, {x0, x1}; \n\t"
+      "mul.f32x2 x01, x01, %2; \n\t"
+      "mov.b64 x23, {x2, x3}; \n\t"
+      "mul.f32x2 x23, x23, %2; \n\t"
+      "mov.b64 {x0, x1}, x01; \n\t"
+      "mov.b64 {x2, x3}, x23; \n\t"
+      "cvt.rn.satfinite.e2m1x2.f32 q0, x1, x0; \n\t"
+      "cvt.rn.satfinite.e2m1x2.f32 q1, x3, x2; \n\t"
+      "mov.b16 %0, {q0, q1}; \n\t"
+      "}"
+      : "=h"(out_fp4x4)
+      : "l"(reinterpret_cast<const uint64_t&>(input_fp16x4)),
+        "l"(reinterpret_cast<const uint64_t&>(scale)));
+  return out_fp4x4;
+}
+
+__device__ __forceinline__ uint16_t scale_cvt_fp16x4_to_fp4x4_rs(
+    const fp16x4 input_fp16x4,
+    const float2 scale,
+    uint32_t rbits) {
+  uint16_t out_fp4x4 = 0;
+  asm volatile(
+      "{\n"
+      ".reg.b16 x0_fp16; \n\t"
+      ".reg.b16 x1_fp16; \n\t"
+      ".reg.b16 x2_fp16; \n\t"
+      ".reg.b16 x3_fp16; \n\t"
+      ".reg.b32 x0; \n\t"
+      ".reg.b32 x1; \n\t"
+      ".reg.b32 x2; \n\t"
+      ".reg.b32 x3; \n\t"
+      ".reg.b64 x01; \n\t"
+      ".reg.b64 x23; \n\t"
+      ".reg.b16 q0; \n\t"
+      "mov.b64 {x0_fp16, x1_fp16, x2_fp16, x3_fp16} , %1; \n\t"
+      "cvt.f32.f16 x0, x0_fp16; \n\t"
+      "cvt.f32.f16 x1, x1_fp16; \n\t"
+      "cvt.f32.f16 x2, x2_fp16; \n\t"
+      "cvt.f32.f16 x3, x3_fp16; \n\t"
+      "mov.b64 x01, {x0, x1}; \n\t"
+      "mul.f32x2 x01, x01, %2; \n\t"
+      "mov.b64 x23, {x2, x3}; \n\t"
+      "mul.f32x2 x23, x23, %2; \n\t"
+      "mov.b64 {x0, x1}, x01; \n\t"
+      "mov.b64 {x2, x3}, x23; \n\t"
+      "cvt.rs.satfinite.e2m1x4.f32 q0, {x3, x2, x1, x0}, %3; \n\t"
+      "}"
+      : "=h"(out_fp4x4)
+      : "l"(reinterpret_cast<const uint64_t&>(input_fp16x4)),
+        "l"(reinterpret_cast<const uint64_t&>(scale)),
+        "r"(rbits));
+  return out_fp4x4;
+}
+
+template <bool USE_SR>
+__device__ __forceinline__ uint16_t scale_cvt_bf16x4_to_fp4x4(
+    const bf16x4 input,
+    const float scale,
+    uint32_t rbits) {
+  float2 scale_fp32x2 = make_float2(scale, scale);
+  if constexpr (USE_SR) {
+    return scale_cvt_bf16x4_to_fp4x4_rs(input, scale_fp32x2, rbits);
+  } else {
+    return scale_cvt_bf16x4_to_fp4x4_rn(input, scale_fp32x2);
+  }
+}
+
+template <bool USE_SR>
+__device__ __forceinline__ uint16_t scale_cvt_fp16x4_to_fp4x4(
+    const fp16x4 input,
+    const float scale,
+    uint32_t rbits) {
+  float2 scale_fp32x2 = make_float2(scale, scale);
+  if constexpr (USE_SR) {
+    return scale_cvt_fp16x4_to_fp4x4_rs(input, scale_fp32x2, rbits);
+  } else {
+    return scale_cvt_fp16x4_to_fp4x4_rn(input, scale_fp32x2);
+  }
+}
+
+template <bool USE_SR>
+__device__ __forceinline__ uint16_t
+scale_cvt_f32x4_to_fp4x4(const f32x4 input, const float scale, uint32_t rbits) {
+  float2 scale_fp32x2 = make_float2(scale, scale);
+  float2 input_fp32x2_0 = make_float2(input.x, input.y);
+  float2 input_fp32x2_1 = make_float2(input.z, input.w);
+
+  if constexpr (USE_SR) {
+    return scale_cvt_fp32x4_to_fp4x4_rs(
+        input_fp32x2_0, input_fp32x2_1, scale_fp32x2, rbits);
+  } else {
+    return scale_cvt_fp32x4_to_fp4x4_rn(
+        input_fp32x2_0, input_fp32x2_1, scale_fp32x2);
+  }
+}
+
+template <typename T, bool USE_SR>
+__device__ __forceinline__ uint16_t scale_cvt_Tx4_to_fp4x4_fast(
+    const Vector4_t<T> input,
+    const float scale,
+    uint32_t rbits) {
+  if constexpr (std::is_same<T, __nv_bfloat16>::value) {
+    return scale_cvt_bf16x4_to_fp4x4<USE_SR>(input, scale, rbits);
+  } else if constexpr (std::is_same<T, __half>::value) {
+    return scale_cvt_fp16x4_to_fp4x4<USE_SR>(input, scale, rbits);
+  } else {
+    return scale_cvt_f32x4_to_fp4x4<USE_SR>(input, scale, rbits);
+  }
+}
+#endif // (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) &&
+       // (__CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000)
+
+template <typename T, bool USE_SR>
+__device__ __forceinline__ uint16_t scale_cvt_Tx4_to_fp4x4(
+    const Vector4_t<T> input,
+    const float scale,
+    uint32_t rbits) {
+#if (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) && \
+    (__CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000)
+  return scale_cvt_Tx4_to_fp4x4_fast<T, USE_SR>(input, scale, rbits);
+#else
+  static_assert(
+      !USE_SR,
+      "Stochastic rounding (USE_SR=true) requires CUDA >= 12.8 and compute capability >= 1000.");
+  return scale_cvt_Tx4_to_fp4x4_fallback(input, scale);
+#endif
+}
+} // namespace mlx::core::cu

mlx/backend/cuda/quantized/qqmm.cpp

@@ -0,0 +1,205 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/quantized/cublas_qqmm.h"
+#include "mlx/backend/cuda/quantized/qqmm_utils.h"
+#include "mlx/backend/cuda/quantized/quantized.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/primitives.h"
+
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+namespace {
+
+inline array ensure_row_contiguous(
+    const array& x,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  if (!x.flags().row_contiguous) {
+    array x_copy = contiguous_copy_gpu(x, s);
+    enc.add_temporary(x_copy);
+    return x_copy;
+  } else {
+    return x;
+  }
+}
+
+inline array ensure_row_contiguous_matrix(
+    const array& x,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  if (x.ndim() < 2) {
+    if (x.strides()[0] == 1) {
+      return x;
+    }
+  } else {
+    auto stride_0 = x.strides()[x.ndim() - 2];
+    auto stride_1 = x.strides()[x.ndim() - 1];
+    if (stride_0 == x.shape(-1) && stride_1 == 1) {
+      return x;
+    }
+  }
+  array x_copy = contiguous_copy_gpu(x, s);
+  enc.add_temporary(x_copy);
+  return x_copy;
+}
+
+array pad_and_repack_scales(
+    const array& scale,
+    cu::CommandEncoder& encoder,
+    const Stream& s) {
+  // Compute padded dimensions for full tiles (128 rows × 4 cols)
+  auto [pad_outer, pad_inner] =
+      get_padded_scale_dims(scale.shape(-2), scale.shape(-1));
+  // cuBLAS requirements for scale factor layout:
+  // 1. Dimensions must be padded to full tiles (128 rows × 4 cols)
+  // 2. Out-of-bounds values must be filled with zeros
+  // 3. Starting addresses must be 16-byte aligned
+  //
+  // https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout
+  // Note: cu::malloc_async already provides 256-byte alignment
+  array scale_tiled(
+      cu::malloc_async(pad_outer * pad_inner, encoder),
+      Shape{pad_outer, pad_inner},
+      scale.dtype());
+  repack_scales(scale, scale_tiled, encoder, s);
+
+  encoder.add_temporary(scale_tiled);
+  return scale_tiled;
+}
+} // namespace
+
+namespace {
+void qqmm_impl(
+    cu::CommandEncoder& encoder,
+    int M,
+    int N,
+    int K,
+    bool a_transposed,
+    int64_t lda,
+    bool b_transposed,
+    int64_t ldb,
+    array& out,
+    const array& a,
+    const array& b,
+    const array& a_scale,
+    const array& b_scale,
+    Dtype out_dtype,
+    QuantizationMode mode,
+    float alpha = 1.0f) {
+  // Invoke CublasQQMM
+  std::string qmode = quantization_mode_to_string(mode);
+
+  // Currently only supports non-batched QQMM operations
+  // that covers all use cases for training, we will just collapse (batch,
+  // seq_len) into (tokens)
+  CublasQQMM qqmm(
+      encoder.device(),
+      a_transposed,
+      M,
+      K,
+      lda,
+      b_transposed,
+      K,
+      N,
+      ldb,
+      1, // batch_count
+      0, // a_batch_stride
+      0, // b_batch_stride
+      out_dtype,
+      qmode);
+
+  qqmm.run(encoder, out, a, b, a_scale, b_scale, alpha);
+}
+} // namespace
+
+void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("QQMatmul::eval_gpu");
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+  auto& device = encoder.device();
+  auto cc = device.compute_capability_major() * 100 +
+      device.compute_capability_minor() * 10;
+  if (cc < 1000) {
+    throw std::runtime_error(
+        "[QQMatmul::eval_gpu] QQMM is only supported on GPUs with compute capability 10.0 or higher.");
+  }
+  assert(
+      (inputs.size() == 3 && inputs[1].dtype() == uint32) ||
+      (inputs.size() == 2));
+
+  auto quantize = [&](const array& input,
+                      cu::CommandEncoder& encoder,
+                      const Stream& s) -> std::pair<array, array> {
+    const array x = ensure_row_contiguous(input, encoder, s);
+
+    auto xq_shape = x.shape();
+    xq_shape.back() = x.shape(-1) * bits_ / 32;
+
+    auto sshape = x.shape();
+    const int64_t scales_inner = x.shape(-1) / group_size_;
+    auto [pad_outer, pad_inner] =
+        get_padded_scale_dims(x.shape(-2), scales_inner);
+    sshape[x.ndim() - 2] = pad_outer;
+    sshape[x.ndim() - 1] = pad_inner;
+    sshape.back() = scales_inner;
+
+    // Allocate outputs
+    const int64_t xq_bytes = x.size() * bits_ / 8;
+    const int64_t batch = x.size() / (x.shape(-2) * x.shape(-1));
+    const int64_t scales_bytes = batch * (pad_outer * pad_inner);
+
+    array x_q(cu::malloc_async(xq_bytes, encoder), std::move(xq_shape), uint32);
+    array scales_x(
+        cu::malloc_async(scales_bytes, encoder), std::move(sshape), uint8);
+
+    fp_quantize(x, x_q, scales_x, group_size_, bits_, encoder, s);
+
+    encoder.add_temporary(x_q);
+    encoder.add_temporary(scales_x);
+    return {x_q, scales_x};
+  };
+  auto [x_q, scale_x_pre] = quantize(inputs[0], encoder, s);
+  auto [w_q, scale_w_pre] = (inputs[1].dtype() != uint32)
+      ? quantize(inputs[1], encoder, s)
+      : std::make_pair(inputs[1], inputs[2]);
+
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+
+  auto out_dtype = out.dtype();
+
+  int M = x_q.shape(-2);
+  int N = w_q.shape(-2); // always transposed
+  int K_packed = x_q.shape(-1);
+  int K = K_packed * (32 / bits_);
+
+  // Repack scales from linear to tiled layout for tensor cores
+  array scale_x = pad_and_repack_scales(scale_x_pre, encoder, s);
+  array scale_w = pad_and_repack_scales(scale_w_pre, encoder, s);
+
+  bool x_transposed = false;
+  bool w_transposed = true; // always transposed
+  int64_t lda = K;
+  int64_t ldb = K;
+
+  qqmm_impl(
+      encoder,
+      M,
+      N,
+      K,
+      x_transposed,
+      lda,
+      w_transposed,
+      ldb,
+      out,
+      x_q,
+      w_q,
+      scale_x,
+      scale_w,
+      out_dtype,
+      mode_);
+}
+
+} // namespace mlx::core

mlx/backend/cuda/quantized/qqmm_utils.cu

@@ -0,0 +1,169 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/cuda/quantized/qqmm_utils.h"
+
+#include <cooperative_groups.h>
+
+namespace mlx::core {
+
+namespace cg = cooperative_groups;
+
+// To pass scales to tensor cores, they need to be repacked into a tiled layout
+// https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout
+// Tiled layout for scale factors is very well described in CUTLASS
+// documentation:
+// https://github.com/NVIDIA/cutlass/blob/main/media/docs/cpp/blackwell_functionality.md#scale-factor-layouts
+// Conceptually, it should be like this:
+// q_w = mx.zeros(shape=(M, N)) <-- zeros just for an example
+// s.shape = (M, N // 16) -- packed in row contigous order, group_size = 16
+// cbg_cnt = N // 16 // 4
+// rb_cnt = M // 128
+// tmp = x.reshape(rb_cnt, 4, 32, cbg_cnt, 4)
+// repacked_scales = tmp.transpose(0, 3, 2, 1, 4)
+// example: indecis of intial tile 128 x 4 of scales (packed in row major tensor
+// (M, K // 16), where M = 128, K = 64): array([[0, 1, 2, 3],
+//       [4, 5, 6, 7],
+//       [8, 9, 10, 11],
+//       ...,
+//       [500, 501, 502, 503],
+//       [504, 505, 506, 507],
+//       [508, 509, 510, 511]]
+// packed scales within tile 128 x 4:
+// array([[[[[0, 1, 2, 3], <-- s_0,0..s_0,3 scales
+//          [128, 129, 130, 131], <-- s_32,0..s_32,3 scales
+//          [256, 257, 258, 259], <-- s_64,0..s_64,3 scales
+//          [384, 385, 386, 387]], <-- s_96,0..s_96,3 scales
+//         [[4, 5, 6, 7], <-- s_1,0..s_1,3 scales
+//          [132, 133, 134, 135], ...
+//          [260, 261, 262, 263],
+//          [388, 389, 390, 391]],
+//         [[124, 125, 126, 127],
+//          [252, 253, 254, 255],
+//          [380, 381, 382, 383],
+//          [508, 509, 510, 511]]]]],
+__device__ size_t
+scale_tiled_offset(size_t scale_index, size_t num_rows, size_t num_scale_cols) {
+  // Compute the tiled layout offset for scale factors used in tensor cores
+  // This function maps from a linear scale index to the tiled layout expected
+  // by tensor cores (and cublaslt).
+  //
+  // Input: linear scale index (e.g., for a matrix M x K with group_size,
+  //        scale_index ranges from 0 to (M * K/group_size - 1))
+  //
+  // The tiled layout organizes scales into tiles of 128 rows x 4 columns,
+  // where each tile is subdivided into 4 sub-blocks of 32 rows x 4 columns.
+  size_t row = scale_index / num_scale_cols;
+  size_t col = scale_index % num_scale_cols;
+
+  constexpr size_t rows_per_tile = 128;
+  constexpr size_t rows_per_sub_block = 32;
+  constexpr size_t cols_per_sub_block = 4;
+  constexpr size_t sub_blocks_per_tile = 4; // Vertically stacked
+
+  // Decompose row position
+  size_t tile_row = row / rows_per_tile; // Which tile row
+  size_t row_in_tile = row % rows_per_tile; // Row within tile
+  size_t sub_block_row =
+      row_in_tile / rows_per_sub_block; // Sub-block within tile
+  size_t row_in_sub_block =
+      row_in_tile % rows_per_sub_block; // Row in sub-block
+
+  // Decompose column position
+  size_t col_tile = col / cols_per_sub_block; // Which column tile
+  size_t col_in_sub_block = col % cols_per_sub_block; // Column within sub-block
+
+  // Compute tile index and offset within tile
+  size_t num_col_tiles = cuda::ceil_div(num_scale_cols, cols_per_sub_block);
+  size_t tile_idx = tile_row * num_col_tiles + col_tile;
+
+  size_t offset_in_tile =
+      (row_in_sub_block * sub_blocks_per_tile * cols_per_sub_block) +
+      (sub_block_row * cols_per_sub_block) + col_in_sub_block;
+
+  constexpr size_t tile_size = rows_per_tile * cols_per_sub_block;
+  return tile_idx * tile_size + offset_in_tile;
+}
+
+namespace cu {
+
+__global__ void repack_scales(
+    const uint8_t* scales_linear,
+    uint8_t* scales_tiled,
+    size_t input_rows,
+    size_t input_cols,
+    size_t output_rows,
+    size_t output_cols) {
+  auto block_size = cg::this_thread_block().dim_threads();
+  auto block_idx = cg::this_thread_block().group_index();
+  auto idx_in_block = cg::this_thread_block().thread_index();
+
+  auto tidx = block_idx.x * block_size.x + idx_in_block.x;
+  auto tidy = block_idx.y * block_size.y + idx_in_block.y;
+
+  auto grid_dim_x =
+      cg::this_grid().dim_blocks().x * cg::this_grid().block_index().x;
+
+  size_t output_index = tidx + grid_dim_x * size_t(tidy);
+  size_t output_size = output_rows * output_cols;
+
+  if (output_index >= output_size) {
+    return;
+  }
+
+  size_t tiled_offset =
+      scale_tiled_offset(output_index, output_rows, output_cols);
+
+  size_t row = output_index / output_cols;
+  size_t col = output_index % output_cols;
+
+  // Probably this can be done better with 2 separated paths for valid and
+  // padding
+  if (row < input_rows && col < input_cols) {
+    size_t input_index = row * input_cols + col;
+    scales_tiled[tiled_offset] = scales_linear[input_index];
+  } else {
+    // Zero-fill padding region
+    scales_tiled[tiled_offset] = 0;
+  }
+}
+
+} // namespace cu
+
+void repack_scales(
+    const array& scales,
+    array& scales_tiled,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  enc.set_input_array(scales);
+  enc.set_output_array(scales_tiled);
+
+  // Note: scales_tiled is padded to full tiles so if num_rows or num_cols
+  // are not multiples of tile sizes, the extra space is filled with zeros
+
+  size_t input_rows = scales.shape(-2);
+  size_t input_cols = scales.shape(-1);
+
+  size_t output_rows = scales_tiled.shape(-2);
+  size_t output_cols = scales_tiled.shape(-1);
+  size_t output_size = output_rows * output_cols;
+
+  bool large = output_size > UINT_MAX;
+  auto [num_blocks, block_dims] = get_launch_args(
+      output_size, scales_tiled.shape(), scales_tiled.strides(), large);
+
+  enc.add_kernel_node(
+      cu::repack_scales,
+      num_blocks,
+      block_dims,
+      0,
+      gpu_ptr<uint8_t>(scales),
+      gpu_ptr<uint8_t>(scales_tiled),
+      input_rows,
+      input_cols,
+      output_rows,
+      output_cols);
+}
+
+} // namespace mlx::core

mlx/backend/cuda/quantized/qqmm_utils.h

@@ -0,0 +1,30 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+#include "mlx/backend/cuda/device.h"
+
+namespace mlx::core {
+
+// Compute padded dimensions for tiled layout
+// Tiles are 128 rows × 4 columns, must allocate full tiles
+inline std::pair<int, int> get_padded_scale_dims(int num_rows, int num_cols) {
+  constexpr int rows_per_tile = 128;
+  constexpr int cols_per_tile = 4;
+
+  int padded_rows =
+      ((num_rows + rows_per_tile - 1) / rows_per_tile) * rows_per_tile;
+  int padded_cols =
+      ((num_cols + cols_per_tile - 1) / cols_per_tile) * cols_per_tile;
+
+  return {padded_rows, padded_cols};
+}
+
+void repack_scales(
+    const array& scales,
+    array& scales_tiled,
+    cu::CommandEncoder& enc,
+    const Stream& s);
+
+} // namespace mlx::core

mlx/backend/cuda/quantized/quantized_utils.cuh

@@ -15,6 +15,22 @@ inline constexpr __device__ short get_bytes_per_pack() {
   return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3);
 }
 
+template <typename T>
+__device__ __forceinline__ void abs_max_x2(T& out, const T& x1, const T& x2) {
+  if constexpr (
+      (std::is_same<T, __nv_bfloat162>::value) ||
+      (std::is_same<T, __half2>::value)) {
+    T a = x1;
+    T b = x2;
+    out = __hmax2(__habs2(a), __habs2(b));
+  } else if constexpr (std::is_same<T, float2>::value) {
+    float2 a = x1;
+    float2 b = x2;
+    out.x = fmaxf(fabsf(a.x), fabsf(b.x));
+    out.y = fmaxf(fabsf(a.y), fabsf(b.y));
+  }
+}
+
 } // namespace cu
 
 template <typename F>

mlx/backend/cuda/random.cu

@@ -139,10 +139,10 @@ void RandomBits::eval_gpu(const std::vector<array>& inputs, array& out) {
   // keys has shape (N1, ..., NK, 2)
   // out has shape (N1, ..., NK, M1, M2, ...)
   auto& keys = inputs[0];
-  uint32_t num_keys = keys.size() / 2;
+  size_t num_keys = keys.size() / 2;
 
-  uint32_t elems_per_key = out.size() / num_keys;
-  uint32_t bytes_per_key = out.itemsize() * elems_per_key;
+  size_t elems_per_key = out.size() / num_keys;
+  size_t bytes_per_key = out.itemsize() * elems_per_key;
   auto& s = stream();
   auto& encoder = cu::get_command_encoder(s);
   out.set_data(cu::malloc_async(out.nbytes(), encoder));
@@ -150,19 +150,25 @@ void RandomBits::eval_gpu(const std::vector<array>& inputs, array& out) {
     return;
   }
 
-  uint32_t out_per_key = (bytes_per_key + 4 - 1) / 4;
-  uint32_t half_size = out_per_key / 2;
+  size_t out_per_key = (bytes_per_key + 4 - 1) / 4;
+  size_t half_size = out_per_key / 2;
+
   bool odd = out_per_key % 2;
+  if ((half_size + odd) >= UINT32_MAX || num_keys >= UINT32_MAX) {
+    throw std::runtime_error("[RandomBits::eval_gpu] Large size unsupported");
+  }
 
   encoder.set_input_array(keys);
   encoder.set_output_array(out);
-  dim3 grid_dims{num_keys, half_size + odd};
-  int64_t total = grid_dims.x * grid_dims.y;
-  int32_t threads_y = 1;
-  while ((total / threads_y) >= (1U << 31)) {
+  int64_t total = num_keys * (half_size + odd);
+  uint32_t threads_y = 1;
+  while ((total / threads_y) >= UINT_MAX) {
     threads_y *= 2;
   }
-  int32_t threads_x = cuda::ceil_div(total, threads_y);
+  uint32_t threads_x = cuda::ceil_div(total, threads_y);
+
+  dim3 grid_dims{
+      static_cast<uint32_t>(num_keys), static_cast<uint32_t>(half_size + odd)};
   auto [grid, block] = get_grid_and_block(threads_x, threads_y, 1);
   auto& stream = encoder.stream();
   if (keys.flags().row_contiguous) {

mlx/backend/cuda/reduce/col_reduce.cu

@@ -89,9 +89,13 @@ template <
     int NDIM,
     int BM,
     int BN,
-    int N_READS = 4>
-__global__ void
-col_reduce_looped(T* in, U* out, const __grid_constant__ ColReduceArgs args) {
+    int N_READS = 4,
+    int BLOCKS = 1>
+__global__ void col_reduce_looped(
+    T* in,
+    U* out,
+    const __grid_constant__ ColReduceArgs args,
+    int64_t out_size) {
   auto grid = cg::this_grid();
   auto block = cg::this_thread_block();
   auto warp = cg::tiled_partition<WARP_SIZE>(block);
@@ -102,6 +106,8 @@ col_reduce_looped(T* in, U* out, const __grid_constant__ ColReduceArgs args) {
   size_t tile_idx = grid.block_rank();
   size_t tile_x = tile_idx % ((args.reduction_stride + BN - 1) / BN);
   size_t tile_y = tile_idx / ((args.reduction_stride + BN - 1) / BN);
+  size_t tile_out = tile_y / out_size;
+  tile_y = tile_y % out_size;
 
   // Compute the indices for the thread within the tile
   short thread_x = block.thread_rank() % threads_per_row;
@@ -118,12 +124,23 @@ col_reduce_looped(T* in, U* out, const __grid_constant__ ColReduceArgs args) {
     totals[i] = ReduceInit<Op, T>::value();
   }
 
-  LoopedElemToLoc<NDIM, (NDIM > 2)> loop(args.reduce_ndim);
-  loop.next(thread_y, args.reduce_shape.data(), args.reduce_strides.data());
   size_t total = args.non_col_reductions * args.reduction_size;
+  size_t per_block, start, end;
+  if constexpr (BLOCKS > 1) {
+    per_block = (total + BLOCKS - 1) / BLOCKS;
+    start = tile_out * per_block + thread_y;
+    end = min((tile_out + 1) * per_block, total);
+  } else {
+    per_block = total;
+    start = thread_y;
+    end = total;
+  }
+
+  LoopedElemToLoc<NDIM, (NDIM > 2)> loop(args.reduce_ndim);
+  loop.next(start, args.reduce_shape.data(), args.reduce_strides.data());
   if (tile_x * BN + BN <= args.reduction_stride) {
     if (args.reduction_stride % N_READS == 0) {
-      for (size_t r = thread_y; r < total; r += BM) {
+      for (size_t r = start; r < end; r += BM) {
         T vals[N_READS];
         cub::LoadDirectBlockedVectorized(thread_x, in + loop.location(), vals);
         for (int i = 0; i < N_READS; i++) {
@@ -132,7 +149,7 @@ col_reduce_looped(T* in, U* out, const __grid_constant__ ColReduceArgs args) {
         loop.next(BM, args.reduce_shape.data(), args.reduce_strides.data());
       }
     } else {
-      for (size_t r = thread_y; r < total; r += BM) {
+      for (size_t r = start; r < end; r += BM) {
         T vals[N_READS];
         cub::LoadDirectBlocked(thread_x, in + loop.location(), vals);
         for (int i = 0; i < N_READS; i++) {
@@ -142,7 +159,7 @@ col_reduce_looped(T* in, U* out, const __grid_constant__ ColReduceArgs args) {
       }
     }
   } else {
-    for (size_t r = thread_y; r < total; r += BM) {
+    for (size_t r = start; r < end; r += BM) {
       T vals[N_READS];
       cub::LoadDirectBlocked(
           thread_x,
@@ -173,6 +190,9 @@ col_reduce_looped(T* in, U* out, const __grid_constant__ ColReduceArgs args) {
 
   // Write result.
   if (warp.thread_rank() == 0) {
+    if (BLOCKS > 1) {
+      out += tile_out * out_size * args.reduction_stride;
+    }
     cub::StoreDirectBlocked(
         warp.meta_group_rank(),
         out + tile_y * args.reduction_stride + tile_x * BN,
@@ -227,11 +247,12 @@ __global__ void col_reduce_small(
 inline auto output_grid_for_col_reduce(
     const array& out,
     const cu::ColReduceArgs& args,
-    int bn) {
+    int bn,
+    int outer = 1) {
   int gx, gy = 1;
   size_t n_inner_blocks = cuda::ceil_div(args.reduction_stride, bn);
   size_t n_outer_blocks = out.size() / args.reduction_stride;
-  size_t n_blocks = n_outer_blocks * n_inner_blocks;
+  size_t n_blocks = n_outer_blocks * n_inner_blocks * outer;
   while (n_blocks / gy > INT32_MAX) {
     gy *= 2;
   }
@@ -277,7 +298,8 @@ void col_reduce_looped(
             0,
             indata,
             gpu_ptr<U>(out),
-            static_cast<cu::ColReduceArgs>(args));
+            static_cast<cu::ColReduceArgs>(args),
+            out.size() / args.reduction_stride);
       });
     });
   });
@@ -320,6 +342,117 @@ void col_reduce_small(
   });
 }
 
+void col_reduce_two_pass(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    array& out,
+    Reduce::ReduceType reduce_type,
+    const std::vector<int>& axes,
+    const ReductionPlan& plan,
+    const cu::ColReduceArgs& args) {
+  // Allocate data for the output using in's layout to access them as
+  // contiguously as possible.
+  allocate_same_layout(out, in, axes, encoder);
+
+  // Allocate an intermediate array to hold the 1st pass result
+  constexpr int outer = 32;
+
+  Shape intermediate_shape;
+  intermediate_shape.push_back(outer);
+  intermediate_shape.insert(
+      intermediate_shape.end(), out.shape().begin(), out.shape().end());
+
+  Strides intermediate_strides;
+  intermediate_strides.push_back(out.size());
+  intermediate_strides.insert(
+      intermediate_strides.end(), out.strides().begin(), out.strides().end());
+
+  array intermediate(intermediate_shape, out.dtype(), nullptr, {});
+  auto [data_size, rc, cc] =
+      check_contiguity(intermediate_shape, intermediate_strides);
+  auto fl = out.flags();
+  fl.row_contiguous = rc;
+  fl.col_contiguous = cc;
+  fl.contiguous = true;
+  intermediate.set_data(
+      cu::malloc_async(intermediate.nbytes(), encoder),
+      data_size,
+      intermediate_strides,
+      fl,
+      allocator::free);
+
+  encoder.add_temporary(intermediate);
+  encoder.set_input_array(in);
+  encoder.set_output_array(intermediate);
+  dispatch_all_types(in.dtype(), [&](auto type_tag) {
+    dispatch_reduce_ops(reduce_type, [&](auto reduce_type_tag) {
+      dispatch_reduce_ndim(args.reduce_ndim, [&](auto reduce_ndim) {
+        using OP = MLX_GET_TYPE(reduce_type_tag);
+        using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+        using U = typename cu::ReduceResult<OP, T>::type;
+        // Cub doesn't like const pointers for vectorized loads. (sigh)
+        T* indata = const_cast<T*>(gpu_ptr<T>(in));
+
+        constexpr int N_READS = 4;
+        constexpr int BM = 32;
+        constexpr int BN = 32;
+        dim3 grid = output_grid_for_col_reduce(out, args, BN, outer);
+        int blocks = BM * BN / N_READS;
+        auto kernel = cu::
+            col_reduce_looped<T, U, OP, reduce_ndim(), BM, BN, N_READS, outer>;
+        encoder.add_kernel_node(
+            kernel,
+            grid,
+            blocks,
+            0,
+            indata,
+            gpu_ptr<U>(intermediate),
+            static_cast<cu::ColReduceArgs>(args),
+            out.size() / args.reduction_stride);
+      });
+    });
+  });
+
+  // Prepare the reduction arguments for the 2nd pass
+  cu::ColReduceArgs second_args = args;
+  second_args.reduction_size = outer;
+  second_args.reduction_stride = out.size();
+  second_args.ndim = 0;
+  second_args.reduce_shape[0] = outer;
+  second_args.reduce_strides[0] = out.size();
+  second_args.reduce_ndim = 1;
+  second_args.non_col_reductions = 1;
+
+  encoder.set_input_array(intermediate);
+  encoder.set_output_array(out);
+  dispatch_all_types(intermediate.dtype(), [&](auto type_tag) {
+    dispatch_reduce_ops(reduce_type, [&](auto reduce_type_tag) {
+      dispatch_reduce_ndim(second_args.reduce_ndim, [&](auto reduce_ndim) {
+        using OP = MLX_GET_TYPE(reduce_type_tag);
+        using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+        using U = typename cu::ReduceResult<OP, T>::type;
+
+        constexpr int N_READS = 4;
+        constexpr int BM = 32;
+        constexpr int BN = 32;
+        dim3 grid = output_grid_for_col_reduce(out, second_args, BN);
+        int blocks = BM * BN / N_READS;
+        auto kernel =
+            cu::col_reduce_looped<T, U, OP, reduce_ndim(), BM, BN, N_READS>;
+        encoder.add_kernel_node(
+            kernel,
+            grid,
+            blocks,
+            0,
+            gpu_ptr<T>(intermediate),
+            gpu_ptr<U>(out),
+            second_args,
+            second_args.reduction_stride);
+      });
+    });
+  });
+}
+
 void col_reduce(
     cu::CommandEncoder& encoder,
     const array& in,
@@ -334,6 +467,18 @@ void col_reduce(
   //   It is a general strided reduce. Each threadblock computes the output for
   //   a subrow of the fast moving axis. For instance 32 elements.
   //
+  // - col_reduce_small
+  //
+  //  It is a column reduce for small columns. Each thread loops over the whole
+  //  column without communicating with any other thread.
+  //
+  // - col_reduce_two_pass
+  //
+  //  It is a reduce for long columns. To increase parallelism, we split the
+  //  reduction in two passes. First we do a column reduce where many
+  //  threadblocks operate on different parts of the reduced axis. Then we
+  //  perform a final column reduce.
+  //
   // Notes: As in row reduce we opt to read as much in order as possible and
   //        leave transpositions as they are (contrary to our Metal backend).
   //
@@ -349,6 +494,14 @@ void col_reduce(
     return;
   }
 
+  // Long column with smallish row
+  size_t total_sums = args.non_col_reductions * args.reduction_size;
+  size_t approx_threads = out.size();
+  if (total_sums / approx_threads > 32) {
+    col_reduce_two_pass(encoder, in, out, reduce_type, axes, plan, args);
+    return;
+  }
+
   // Fallback col reduce
   col_reduce_looped(encoder, in, out, reduce_type, axes, plan, args);
 }

mlx/backend/cuda/rms_norm.cu

@@ -22,26 +22,28 @@ inline __device__ float2 plus_f2(const float2& a, const float2& b) {
 }
 
 // Similar to cub::BlockReduce, but result is broadcasted to every thread.
-template <typename T, int BLOCK_DIM>
+template <typename T, int BLOCK_DIM, int GROUP_DIM = WARP_SIZE>
 struct BlockBroadcastReduce {
-  static_assert(WARP_SIZE <= BLOCK_DIM && BLOCK_DIM <= WARP_SIZE * WARP_SIZE);
-  static_assert(BLOCK_DIM % WARP_SIZE == 0);
-  using TempStorage = T[BLOCK_DIM / WARP_SIZE];
+  using TempStorage = T[std::max(BLOCK_DIM / WARP_SIZE, 1)];
 
   cg::thread_block& block;
   TempStorage& temp;
 
   template <typename Op>
   __device__ T Reduce(const T& input, const Op& op, const T& init_value) {
-    auto warp = cg::tiled_partition<WARP_SIZE>(block);
+    auto warp = cg::tiled_partition<GROUP_DIM>(block);
     T x = cg::reduce(warp, input, op);
-    if (warp.thread_rank() == 0) {
-      temp[warp.meta_group_rank()] = x;
+    if constexpr (BLOCK_DIM > GROUP_DIM) {
+      if (warp.thread_rank() == 0) {
+        temp[warp.meta_group_rank()] = x;
+      }
+      block.sync();
+      x = warp.thread_rank() < warp.meta_group_size() ? temp[warp.thread_rank()]
+                                                      : init_value;
+      return cg::reduce(warp, x, op);
+    } else {
+      return x;
     }
-    block.sync();
-    x = warp.thread_rank() < warp.meta_group_size() ? temp[warp.thread_rank()]
-                                                    : init_value;
-    return cg::reduce(warp, x, op);
   }
 
   __device__ T Sum(const T& input) {
@@ -49,6 +51,52 @@ struct BlockBroadcastReduce {
   }
 };
 
+template <typename T, int BLOCK_DIM, int REDUCE_DIM, int N_READS = 4>
+__global__ void rms_norm_small(
+    const T* x,
+    const T* w,
+    T* out,
+    float eps,
+    uint32_t axis_size,
+    uint32_t n_rows,
+    int64_t w_stride) {
+  auto grid = cg::this_grid();
+  auto block = cg::this_thread_block();
+
+  using BlockReduceT = BlockBroadcastReduce<float, BLOCK_DIM, REDUCE_DIM>;
+  __shared__ typename BlockReduceT::TempStorage temp;
+
+  auto row =
+      (grid.block_rank() * block.dim_threads().y) + block.thread_index().y;
+  if (row >= n_rows) {
+    return;
+  }
+  x += row * axis_size;
+  out += row * axis_size;
+
+  // Normalizer.
+  float normalizer = 0;
+  auto index = block.thread_index().x;
+  auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    float t = static_cast<float>(xn[i]);
+    normalizer += t * t;
+  }
+
+  normalizer = BlockReduceT{block, temp}.Sum(normalizer);
+  normalizer = rsqrt(normalizer / axis_size + eps);
+
+  // Outputs.
+  auto wn = load_vector<N_READS>(w, index, axis_size, w_stride, T(0));
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    float y = static_cast<float>(xn[i]) * normalizer;
+    xn[i] = wn[i] * static_cast<T>(y);
+  }
+  store_vector<N_READS>(out, index, xn, axis_size);
+}
+
 template <typename T, int BLOCK_DIM, int N_READS = 4>
 __global__ void rms_norm(
     const T* x,
@@ -94,6 +142,74 @@ __global__ void rms_norm(
   }
 }
 
+template <
+    typename T,
+    bool HAS_W,
+    int BLOCK_DIM,
+    int REDUCE_DIM,
+    int N_READS = 4>
+__global__ void rms_norm_vjp_small(
+    const T* x,
+    const T* w,
+    const T* g,
+    T* gx,
+    T* gw,
+    float eps,
+    int32_t axis_size,
+    int32_t n_rows,
+    int64_t w_stride) {
+  auto grid = cg::this_grid();
+  auto block = cg::this_thread_block();
+
+  using BlockReduceF2 = BlockBroadcastReduce<float2, BLOCK_DIM, REDUCE_DIM>;
+  __shared__ typename BlockReduceF2::TempStorage temp;
+
+  auto row =
+      (grid.block_rank() * block.dim_threads().y) + block.thread_index().y;
+  if (row >= n_rows) {
+    return;
+  }
+
+  x += row * axis_size;
+  g += row * axis_size;
+  gx += row * axis_size;
+  gw += row * axis_size;
+
+  // Normalizer.
+  float2 factors = {};
+  auto index = block.thread_index().x;
+  auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+  auto gn = load_vector<N_READS>(g, index, axis_size, T(0));
+  auto wn = load_vector<N_READS>(w, index, axis_size, w_stride, T(0));
+  for (int i = 0; i < N_READS; i++) {
+    float t = static_cast<float>(xn[i]);
+    float wi = wn[i];
+    float gi = gn[i];
+    float wg = wi * gi;
+    factors = plus_f2(factors, {wg * t, t * t});
+  }
+
+  factors = BlockReduceF2{block, temp}.Reduce(factors, plus_f2, {});
+  float meangwx = factors.x / axis_size;
+  float normalizer = rsqrt(factors.y / axis_size + eps);
+  float normalizer3 = normalizer * normalizer * normalizer;
+
+  // Outputs.
+  for (int i = 0; i < N_READS; i++) {
+    float xi = xn[i];
+    float wi = wn[i];
+    float gi = gn[i];
+    xn[i] = static_cast<T>(normalizer * wi * gi - xi * meangwx * normalizer3);
+    if constexpr (HAS_W) {
+      wn[i] = static_cast<T>(gi * xi * normalizer);
+    }
+  }
+  store_vector<N_READS>(gx, index, xn, axis_size);
+  if constexpr (HAS_W) {
+    store_vector<N_READS>(gw, index, wn, axis_size);
+  }
+}
+
 template <typename T, bool HAS_W, int BLOCK_DIM, int N_READS = 4>
 __global__ void rms_norm_vjp(
     const T* x,
@@ -107,12 +223,8 @@ __global__ void rms_norm_vjp(
   auto grid = cg::this_grid();
   auto block = cg::this_thread_block();
 
-  using BlockReduceF = BlockBroadcastReduce<float, BLOCK_DIM>;
   using BlockReduceF2 = BlockBroadcastReduce<float2, BLOCK_DIM>;
-  __shared__ union {
-    typename BlockReduceF::TempStorage f;
-    typename BlockReduceF2::TempStorage f2;
-  } temp;
+  __shared__ typename BlockReduceF2::TempStorage temp;
 
   x += grid.block_rank() * axis_size;
   g += grid.block_rank() * axis_size;
@@ -134,7 +246,7 @@ __global__ void rms_norm_vjp(
       factors = plus_f2(factors, {wg * t, t * t});
     }
   }
-  factors = BlockReduceF2{block, temp.f2}.Reduce(factors, plus_f2, {});
+  factors = BlockReduceF2{block, temp}.Reduce(factors, plus_f2, {});
   float meangwx = factors.x / axis_size;
   float normalizer = rsqrt(factors.y / axis_size + eps);
   float normalizer3 = normalizer * normalizer * normalizer;
@@ -169,6 +281,43 @@ bool RMSNorm::use_fallback(Stream s) {
   return s.device == Device::cpu;
 }
 
+template <int n_per_thread, typename F>
+void dispatch_group_dim(int axis_size, F&& f) {
+  if (axis_size <= n_per_thread * 8) {
+    f(std::integral_constant<int, 8>{},
+      std::integral_constant<int, 1>(),
+      std::integral_constant<int, 16>());
+  } else if (axis_size <= n_per_thread * 16) {
+    f(std::integral_constant<int, 16>{},
+      std::integral_constant<int, 1>(),
+      std::integral_constant<int, 8>());
+  } else if (axis_size <= n_per_thread * 32) {
+    f(std::integral_constant<int, 32>{},
+      std::integral_constant<int, 1>(),
+      std::integral_constant<int, 4>());
+  } else if (axis_size <= n_per_thread * 32 * 2) {
+    f(std::integral_constant<int, 32>{},
+      std::integral_constant<int, 2>(),
+      std::integral_constant<int, 2>());
+  } else if (axis_size <= n_per_thread * 32 * 4) {
+    f(std::integral_constant<int, 32>{},
+      std::integral_constant<int, 4>(),
+      std::integral_constant<int, 1>());
+  } else if (axis_size <= n_per_thread * 32 * 8) {
+    f(std::integral_constant<int, 32>{},
+      std::integral_constant<int, 8>(),
+      std::integral_constant<int, 1>());
+  } else if (axis_size <= n_per_thread * 32 * 16) {
+    f(std::integral_constant<int, 32>{},
+      std::integral_constant<int, 16>(),
+      std::integral_constant<int, 1>());
+  } else {
+    f(std::integral_constant<int, 32>{},
+      std::integral_constant<int, 32>(),
+      std::integral_constant<int, 1>());
+  }
+}
+
 // TODO: There are duplicate code with backend/metal/normalization.cpp
 void RMSNorm::eval_gpu(
     const std::vector<array>& inputs,
@@ -216,12 +365,33 @@ void RMSNorm::eval_gpu(
   dispatch_float_types(out.dtype(), "rms_norm", [&](auto type_tag) {
     using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
     constexpr int N_READS = 16 / sizeof(DataType);
-    dispatch_block_dim(cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) {
-      auto kernel = cu::rms_norm<DataType, block_dim(), N_READS>;
+    if (axis_size <= N_READS * 1024) {
+      dispatch_group_dim<N_READS>(
+          axis_size, [&](auto group_dim, auto n_groups, auto groups_per_block) {
+            constexpr int block_dim = n_groups() * group_dim();
+            auto kernel =
+                cu::rms_norm_small<DataType, block_dim, group_dim(), N_READS>;
+            auto n_blocks =
+                (n_rows + groups_per_block() - 1) / groups_per_block();
+            encoder.add_kernel_node(
+                kernel,
+                n_blocks,
+                {block_dim, groups_per_block()},
+                0,
+                gpu_ptr<DataType>(x),
+                gpu_ptr<DataType>(w),
+                gpu_ptr<DataType>(out),
+                eps_,
+                axis_size,
+                n_rows,
+                w_stride);
+          });
+    } else {
+      auto kernel = cu::rms_norm<DataType, 1024, N_READS>;
       encoder.add_kernel_node(
           kernel,
           n_rows,
-          block_dim(),
+          1024,
           0,
           gpu_ptr<DataType>(x),
           gpu_ptr<DataType>(w),
@@ -229,7 +399,7 @@ void RMSNorm::eval_gpu(
           eps_,
           axis_size,
           w_stride);
-    });
+    }
   });
 }
 
@@ -306,27 +476,51 @@ void RMSNormVJP::eval_gpu(
     dispatch_bool(has_w, [&](auto has_w_constant) {
       using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
       constexpr int N_READS = 16 / sizeof(DataType);
-      dispatch_block_dim(
-          cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) {
-            auto kernel = cu::rms_norm_vjp<
-                DataType,
-                has_w_constant.value,
-                block_dim(),
-                N_READS>;
-            encoder.add_kernel_node(
-                kernel,
-                n_rows,
-                block_dim(),
-                0,
-                gpu_ptr<DataType>(x),
-                gpu_ptr<DataType>(w),
-                gpu_ptr<DataType>(g),
-                gpu_ptr<DataType>(gx),
-                gpu_ptr<DataType>(gw_temp),
-                eps_,
-                axis_size,
-                w_stride);
-          });
+      if (axis_size <= N_READS * 1024) {
+        dispatch_group_dim<N_READS>(
+            axis_size,
+            [&](auto group_dim, auto n_groups, auto groups_per_block) {
+              constexpr int block_dim = group_dim() * n_groups();
+              auto kernel = cu::rms_norm_vjp_small<
+                  DataType,
+                  has_w_constant.value,
+                  block_dim,
+                  group_dim(),
+                  N_READS>;
+              auto n_blocks =
+                  (n_rows + groups_per_block() - 1) / groups_per_block();
+              encoder.add_kernel_node(
+                  kernel,
+                  n_blocks,
+                  {block_dim, groups_per_block()},
+                  0,
+                  gpu_ptr<DataType>(x),
+                  gpu_ptr<DataType>(w),
+                  gpu_ptr<DataType>(g),
+                  gpu_ptr<DataType>(gx),
+                  gpu_ptr<DataType>(gw_temp),
+                  eps_,
+                  axis_size,
+                  n_rows,
+                  w_stride);
+            });
+      } else {
+        auto kernel =
+            cu::rms_norm_vjp<DataType, has_w_constant.value, 1024, N_READS>;
+        encoder.add_kernel_node(
+            kernel,
+            n_rows,
+            1024,
+            0,
+            gpu_ptr<DataType>(x),
+            gpu_ptr<DataType>(w),
+            gpu_ptr<DataType>(g),
+            gpu_ptr<DataType>(gx),
+            gpu_ptr<DataType>(gw_temp),
+            eps_,
+            axis_size,
+            w_stride);
+      }
     });
   });
 

mlx/backend/cuda/scaled_dot_product_attention.cpp

@@ -10,46 +10,8 @@
 
 namespace mlx::core {
 
-namespace fe = cudnn_frontend;
-
 namespace {
 
-#define CHECK_CUDNN_FE_ERROR(cmd)                                    \
-  do {                                                               \
-    auto error = cmd;                                                \
-    if (!error.is_good()) {                                          \
-      throw std::runtime_error(                                      \
-          fmt::format("{} failed: {}.", #cmd, error.get_message())); \
-    }                                                                \
-  } while (0)
-
-std::vector<int64_t> normalized_strides(const array& x) {
-  std::vector<int64_t> strides(x.strides().begin(), x.strides().end());
-  if (std::all_of(
-          strides.begin(), strides.end(), [](int64_t s) { return s == 0; })) {
-    strides.back() = 1;
-    return strides;
-  }
-  if (!x.flags().row_contiguous || x.ndim() < 2) {
-    return strides;
-  }
-  for (int i = x.ndim() - 2; i >= 0; --i) {
-    if (x.shape(i) == 1) {
-      strides[i] = x.shape(i + 1) * strides[i + 1];
-    }
-  }
-  return strides;
-}
-
-void set_tensor_attrs(
-    std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
-    int64_t uid,
-    const array& x) {
-  tensor->set_uid(uid)
-      .set_dim({x.shape().begin(), x.shape().end()})
-      .set_stride(normalized_strides(x));
-}
-
 array prepare_sdpa_input(const array& x, Stream s) {
   // SDPA kernel's requirements on inputs:
   // 1. last dim's stride be 1;
@@ -63,11 +25,43 @@ array prepare_sdpa_input(const array& x, Stream s) {
   return x;
 }
 
+void malloc_with_same_layout(
+    cu::CommandEncoder& encoder,
+    array& o,
+    const array& q) {
+  if (q.flags().row_contiguous) {
+    o.set_data(cu::malloc_async(o.nbytes(), encoder));
+    return;
+  }
+  // fill_order = argsort(q.strides())
+  Shape fill_order(q.ndim());
+  std::iota(fill_order.begin(), fill_order.end(), 0);
+  std::stable_sort(
+      fill_order.begin(), fill_order.end(), [&q](int idx1, int idx2) {
+        auto s1 = q.strides(idx1) > 0 ? q.strides(idx1) : 1;
+        auto s2 = q.strides(idx2) > 0 ? q.strides(idx2) : 1;
+        return s1 < s2;
+      });
+  // Generate o_strides with fill_order
+  Strides o_strides(q.ndim());
+  int64_t stride = 1;
+  for (int i : fill_order) {
+    o_strides[i] = stride;
+    stride *= o.shape(i);
+  }
+  // o is a transposed contiguous array
+  o.set_data(
+      cu::malloc_async(o.nbytes(), encoder),
+      o.size(),
+      o_strides,
+      {true, false, false});
+}
+
 constexpr int QKV_NDIM = 4;
 
 struct SDPACacheKey {
   int device_id;
-  cudnnDataType_t cudnn_dtype;
+  fe::DataType_t cudnn_dtype;
   std::array<int, QKV_NDIM> q_shape;
   std::array<int, QKV_NDIM> k_shape;
   std::array<int, QKV_NDIM> v_shape;
@@ -75,6 +69,8 @@ struct SDPACacheKey {
   std::array<int64_t, QKV_NDIM> k_strides;
   std::array<int64_t, QKV_NDIM> v_strides;
   bool do_causal;
+  std::array<int, QKV_NDIM> mask_shape;
+  std::array<int64_t, QKV_NDIM> mask_strides;
   bool output_logsumexp;
 };
 
@@ -84,6 +80,7 @@ inline BytesKey<SDPACacheKey> build_sdpa_cache_key(
     const array& k,
     const array& v,
     bool do_causal,
+    const std::optional<array>& mask_arr,
     bool output_logsumexp = true) {
   BytesKey<SDPACacheKey> cache_key;
   cache_key.pod = {
@@ -96,20 +93,26 @@ inline BytesKey<SDPACacheKey> build_sdpa_cache_key(
       vector_key<QKV_NDIM>(k.strides()),
       vector_key<QKV_NDIM>(v.strides()),
       do_causal,
+      {},
+      {},
       output_logsumexp,
   };
+  if (mask_arr) {
+    cache_key.pod.mask_shape = vector_key<QKV_NDIM>(mask_arr->shape());
+    cache_key.pod.mask_strides = vector_key<QKV_NDIM>(mask_arr->strides());
+  }
   return cache_key;
 }
 
 auto& sdpa_cache() {
-  static LRUBytesKeyCache<SDPACacheKey, fe::graph::Graph> cache(
-      "MLX_CUDA_SDPA_CACHE_SIZE", /* default_capacity */ 16);
+  static LRUBytesKeyCache<SDPACacheKey, DnnGraph> cache(
+      "MLX_CUDA_SDPA_CACHE_SIZE", /* default_capacity */ 64);
   return cache;
 }
 
 auto& sdpa_backward_cache() {
-  static LRUBytesKeyCache<SDPACacheKey, fe::graph::Graph> cache(
-      "MLX_CUDA_SDPA_BACKWARD_CACHE_SIZE", /* default_capacity */ 16);
+  static LRUBytesKeyCache<SDPACacheKey, DnnGraph> cache(
+      "MLX_CUDA_SDPA_BACKWARD_CACHE_SIZE", /* default_capacity */ 64);
   return cache;
 }
 
@@ -118,6 +121,7 @@ enum UIDS {
   K,
   V,
   SCALE,
+  BIAS,
   O,
   STATS,
   // Backward graph:
@@ -127,166 +131,105 @@ enum UIDS {
   D_O,
 };
 
-fe::graph::Graph build_sdpa_graph(
+DnnGraph build_sdpa_graph(
     cudnnHandle_t handle,
     const array& q,
     const array& k,
     const array& v,
     bool do_causal,
+    const std::optional<array>& mask_arr,
     bool output_logsumexp,
     const array& o,
     const array& stats) {
-  auto dtype = fe::DataType_t::HALF;
-  if (q.dtype() == bfloat16) {
-    dtype = fe::DataType_t::BFLOAT16;
-  }
+  DnnGraph graph(handle, q.dtype());
 
-  fe::graph::Graph graph;
-  graph.set_io_data_type(dtype)
-      .set_intermediate_data_type(fe::DataType_t::FLOAT)
-      .set_compute_data_type(fe::DataType_t::FLOAT);
-
-  auto q_ = graph.tensor(fe::graph::Tensor_attributes().set_name("Q"));
-  auto k_ = graph.tensor(fe::graph::Tensor_attributes().set_name("K"));
-  auto v_ = graph.tensor(fe::graph::Tensor_attributes().set_name("V"));
-  set_tensor_attrs(q_, Q, q);
-  set_tensor_attrs(k_, K, k);
-  set_tensor_attrs(v_, V, v);
-
-  auto scale = graph.tensor(fe::graph::Tensor_attributes()
-                                .set_name("Scale")
-                                .set_uid(SCALE)
-                                .set_dim({1, 1, 1, 1})
-                                .set_stride({1, 1, 1, 1})
-                                .set_is_pass_by_value(true)
-                                .set_data_type(fe::DataType_t::FLOAT));
+  auto q_ = graph.tensor("Q", Q, q);
+  auto k_ = graph.tensor("K", K, k);
+  auto v_ = graph.tensor("V", V, v);
 
   auto options = fe::graph::SDPA_attributes()
                      .set_name("sdpa_cudnn")
-                     .set_attn_scale(scale)
-                     .set_causal_mask(do_causal)
+                     .set_attn_scale(graph.scalar("Scale", SCALE, float32))
                      .set_generate_stats(output_logsumexp);
-
-  auto [o_, stats_] = graph.sdpa(q_, k_, v_, options);
-  o_->set_output(true);
-  set_tensor_attrs(o_, O, o);
-  if (output_logsumexp) {
-    stats_->set_output(true).set_data_type(fe::DataType_t::FLOAT);
-    set_tensor_attrs(stats_, STATS, stats);
+  if (do_causal) {
+    if (q.shape(2) > k.shape(2)) {
+      options.set_causal_mask(do_causal);
+    } else {
+      options.set_causal_mask_bottom_right(do_causal);
+    }
+  }
+  if (mask_arr) {
+    options.set_bias(graph.tensor("BIAS", BIAS, *mask_arr));
   }
 
-  CHECK_CUDNN_FE_ERROR(graph.validate());
-  CHECK_CUDNN_FE_ERROR(graph.build_operation_graph(handle));
-  CHECK_CUDNN_FE_ERROR(graph.create_execution_plans({fe::HeurMode_t::A}));
+  auto [o_, stats_] = graph.sdpa(q_, k_, v_, options);
+  graph.tensor(o_, O, o)->set_output(true);
+  if (output_logsumexp) {
+    graph.tensor(stats_, STATS, stats)->set_output(true);
+  }
+
+  CHECK_CUDNN_FE_ERROR(graph.prepare());
   graph.select_behavior_notes(
       {fe::BehaviorNote_t::SUPPORTS_CUDA_GRAPH_NATIVE_API});
-  CHECK_CUDNN_FE_ERROR(graph.check_support(handle));
-  CHECK_CUDNN_FE_ERROR(graph.build_plans(handle));
-
+  CHECK_CUDNN_FE_ERROR(graph.build());
   return graph;
 }
 
-fe::graph::Graph build_sdpa_backward_graph(
+DnnGraph build_sdpa_backward_graph(
     cudnnHandle_t handle,
     const array& q,
     const array& k,
     const array& v,
     bool do_causal,
+    const std::optional<array>& mask_arr,
     const array& o,
     const array& d_o,
     const array& stats,
     array& d_q,
     array& d_k,
     array& d_v) {
-  auto dtype = fe::DataType_t::HALF;
-  if (q.dtype() == bfloat16) {
-    dtype = fe::DataType_t::BFLOAT16;
-  }
+  DnnGraph graph(handle, q.dtype());
 
-  fe::graph::Graph graph;
-  graph.set_io_data_type(dtype)
-      .set_intermediate_data_type(fe::DataType_t::FLOAT)
-      .set_compute_data_type(fe::DataType_t::FLOAT);
-
-  auto q_ = graph.tensor(fe::graph::Tensor_attributes().set_name("Q"));
-  auto k_ = graph.tensor(fe::graph::Tensor_attributes().set_name("K"));
-  auto v_ = graph.tensor(fe::graph::Tensor_attributes().set_name("V"));
-  auto o_ = graph.tensor(fe::graph::Tensor_attributes().set_name("O"));
-  auto d_o_ = graph.tensor(fe::graph::Tensor_attributes().set_name("D_O"));
-  auto stats_ = graph.tensor(fe::graph::Tensor_attributes().set_name("STATS"));
-  set_tensor_attrs(q_, Q, q);
-  set_tensor_attrs(k_, K, k);
-  set_tensor_attrs(v_, V, v);
-  set_tensor_attrs(o_, O, o);
-  set_tensor_attrs(d_o_, D_O, d_o);
-  set_tensor_attrs(stats_, STATS, stats);
-  stats_->set_data_type(fe::DataType_t::FLOAT);
-
-  auto scale = graph.tensor(fe::graph::Tensor_attributes()
-                                .set_name("Scale")
-                                .set_uid(SCALE)
-                                .set_dim({1, 1, 1, 1})
-                                .set_stride({1, 1, 1, 1})
-                                .set_is_pass_by_value(true)
-                                .set_data_type(fe::DataType_t::FLOAT));
+  auto q_ = graph.tensor("Q", Q, q);
+  auto k_ = graph.tensor("K", K, k);
+  auto v_ = graph.tensor("V", V, v);
+  auto o_ = graph.tensor("O", O, o);
+  auto d_o_ = graph.tensor("D_O", D_O, d_o);
+  auto stats_ = graph.tensor("STATS", STATS, stats);
 
   auto options = fe::graph::SDPA_backward_attributes()
                      .set_name("sdpa_backward_cudnn")
-                     .set_attn_scale(scale)
-                     .set_causal_mask(do_causal);
+                     .set_attn_scale(graph.scalar("Scale", SCALE, float32));
+  if (do_causal) {
+    if (q.shape(2) > k.shape(2)) {
+      options.set_causal_mask(do_causal);
+    } else {
+      options.set_causal_mask_bottom_right(do_causal);
+    }
+  }
+  if (mask_arr) {
+    options.set_bias(graph.tensor("BIAS", BIAS, *mask_arr));
+  }
 
   auto [d_q_, d_k_, d_v_] =
       graph.sdpa_backward(q_, k_, v_, o_, d_o_, stats_, options);
-  d_q_->set_output(true);
-  d_k_->set_output(true);
-  d_v_->set_output(true);
-  set_tensor_attrs(d_q_, D_Q, d_q);
-  set_tensor_attrs(d_k_, D_K, d_k);
-  set_tensor_attrs(d_v_, D_V, d_v);
+  graph.tensor(d_q_, D_Q, d_q)->set_output(true);
+  graph.tensor(d_k_, D_K, d_k)->set_output(true);
+  graph.tensor(d_v_, D_V, d_v)->set_output(true);
 
-  CHECK_CUDNN_FE_ERROR(graph.validate());
-  CHECK_CUDNN_FE_ERROR(graph.build_operation_graph(handle));
-  CHECK_CUDNN_FE_ERROR(graph.create_execution_plans({fe::HeurMode_t::A}));
+  CHECK_CUDNN_FE_ERROR(graph.prepare());
   graph.select_behavior_notes(
       {fe::BehaviorNote_t::SUPPORTS_CUDA_GRAPH_NATIVE_API});
-  CHECK_CUDNN_FE_ERROR(graph.check_support(handle));
-  CHECK_CUDNN_FE_ERROR(graph.build_plans(handle));
-
+  CHECK_CUDNN_FE_ERROR(graph.build());
   return graph;
 }
 
-void execute_graph(
-    cu::CommandEncoder& encoder,
-    cudnnHandle_t handle,
-    fe::graph::Graph& graph,
-    std::unordered_map<int64_t, void*>& variant_pack) {
-  int64_t workspace_size = 0;
-  CHECK_CUDNN_FE_ERROR(graph.get_workspace_size(workspace_size));
-  void* workspace_ptr = nullptr;
-  if (workspace_size > 0) {
-    array workspace(
-        cu::malloc_async(workspace_size, encoder),
-        {static_cast<int>(workspace_size)},
-        uint8);
-    encoder.add_temporary(workspace);
-    workspace_ptr = gpu_ptr<void>(workspace);
-  }
-
-  cudnnSetStream(handle, encoder.stream());
-
-  CudaGraph cuda_graph(encoder.device());
-  CHECK_CUDNN_FE_ERROR(graph.populate_cuda_graph(
-      handle, variant_pack, workspace_ptr, cuda_graph));
-  encoder.add_graph_node(cuda_graph);
-}
-
 } // namespace
 
 bool supports_sdpa_cudnn(
     const array& q,
     const array& k,
     const array& v,
-    bool has_mask,
     bool do_causal,
     Stream s) {
   static bool enabled = env::get_var("MLX_CUDA_USE_CUDNN_SPDA", 1);
@@ -299,19 +242,8 @@ bool supports_sdpa_cudnn(
     return false;
   }
 
-  if (has_mask) {
-    // TODO: Support array masks.
-    if (!do_causal) {
-      return false;
-    }
-    // FIXME: Causal mask generates wrong results when L_Q != L_K.
-    if (q.shape(2) != k.shape(2)) {
-      return false;
-    }
-  }
-
-  // Only use cuDNN for prefilling and training.
-  if (q.shape(2) != k.shape(2)) {
+  // Only use cuDNN for prefilling (T_q > 1) and training (T_q == T_kv).
+  if ((q.shape(2) == 1) && (q.shape(2) != k.shape(2))) {
     return false;
   }
 
@@ -333,47 +265,51 @@ void sdpa_cudnn(
     array& o,
     array& stats,
     bool do_causal,
+    const std::optional<array>& mask_arr,
     bool output_logsumexp,
     Stream s) {
   auto& encoder = cu::get_command_encoder(s);
   auto handle = encoder.device().cudnn_handle();
 
-  // TODO: Handle donation.
-  // TODO: Make O use same memory layout with Q.
-  o.set_data(cu::malloc_async(o.nbytes(), encoder));
+  malloc_with_same_layout(encoder, o, q);
 
   encoder.set_input_array(q);
   encoder.set_input_array(k);
   encoder.set_input_array(v);
   encoder.set_output_array(o);
-
+  if (mask_arr) {
+    encoder.set_input_array(*mask_arr);
+  }
   if (output_logsumexp) {
     stats.set_data(cu::malloc_async(stats.nbytes(), encoder));
     encoder.set_output_array(stats);
   }
 
   // Search cache.
-  auto cache_key =
-      build_sdpa_cache_key(encoder, q, k, v, do_causal, output_logsumexp);
+  auto cache_key = build_sdpa_cache_key(
+      encoder, q, k, v, do_causal, mask_arr, output_logsumexp);
   auto it = sdpa_cache().find(cache_key);
   if (it == sdpa_cache().end()) {
     auto graph = build_sdpa_graph(
-        handle, q, k, v, do_causal, output_logsumexp, o, stats);
+        handle, q, k, v, do_causal, mask_arr, output_logsumexp, o, stats);
     it = sdpa_cache().emplace(cache_key, std::move(graph)).first;
   }
   auto& graph = it->second;
 
   std::unordered_map<int64_t, void*> variant_pack{
-      {Q, const_cast<void*>(gpu_ptr<void>(q))},
-      {K, const_cast<void*>(gpu_ptr<void>(k))},
-      {V, const_cast<void*>(gpu_ptr<void>(v))},
+      {Q, gpu_ptr<void>(q)},
+      {K, gpu_ptr<void>(k)},
+      {V, gpu_ptr<void>(v)},
       {SCALE, &scale},
       {O, gpu_ptr<void>(o)}};
+  if (mask_arr) {
+    variant_pack[BIAS] = gpu_ptr<void>(*mask_arr);
+  }
   if (output_logsumexp) {
     variant_pack[STATS] = gpu_ptr<void>(stats);
   }
 
-  execute_graph(encoder, handle, graph, variant_pack);
+  CHECK_CUDNN_FE_ERROR(graph.encode_graph(encoder, std::move(variant_pack)));
 }
 
 void sdpa_backward_cudnn(
@@ -384,6 +320,7 @@ void sdpa_backward_cudnn(
     const array& o,
     const array& stats,
     bool do_causal,
+    const std::optional<array>& mask_arr,
     const array& d_o,
     array& d_q,
     array& d_k,
@@ -392,10 +329,9 @@ void sdpa_backward_cudnn(
   auto& encoder = cu::get_command_encoder(s);
   auto handle = encoder.device().cudnn_handle();
 
-  // TODO: Handle donation.
-  d_q.set_data(cu::malloc_async(d_q.nbytes(), encoder));
-  d_k.set_data(cu::malloc_async(d_k.nbytes(), encoder));
-  d_v.set_data(cu::malloc_async(d_v.nbytes(), encoder));
+  malloc_with_same_layout(encoder, d_q, q);
+  malloc_with_same_layout(encoder, d_k, k);
+  malloc_with_same_layout(encoder, d_v, v);
 
   encoder.set_input_array(q);
   encoder.set_input_array(k);
@@ -406,30 +342,36 @@ void sdpa_backward_cudnn(
   encoder.set_output_array(d_q);
   encoder.set_output_array(d_k);
   encoder.set_output_array(d_v);
+  if (mask_arr) {
+    encoder.set_input_array(*mask_arr);
+  }
 
   // Search cache.
-  auto cache_key = build_sdpa_cache_key(encoder, q, k, v, do_causal);
+  auto cache_key = build_sdpa_cache_key(encoder, q, k, v, do_causal, mask_arr);
   auto it = sdpa_backward_cache().find(cache_key);
   if (it == sdpa_backward_cache().end()) {
     auto graph = build_sdpa_backward_graph(
-        handle, q, k, v, do_causal, o, d_o, stats, d_q, d_k, d_v);
+        handle, q, k, v, do_causal, mask_arr, o, d_o, stats, d_q, d_k, d_v);
     it = sdpa_backward_cache().emplace(cache_key, std::move(graph)).first;
   }
   auto& graph = it->second;
 
   std::unordered_map<int64_t, void*> variant_pack{
-      {Q, const_cast<void*>(gpu_ptr<void>(q))},
-      {K, const_cast<void*>(gpu_ptr<void>(k))},
-      {V, const_cast<void*>(gpu_ptr<void>(v))},
+      {Q, gpu_ptr<void>(q)},
+      {K, gpu_ptr<void>(k)},
+      {V, gpu_ptr<void>(v)},
       {SCALE, &scale},
-      {O, const_cast<void*>(gpu_ptr<void>(o))},
-      {STATS, const_cast<void*>(gpu_ptr<void>(stats))},
-      {D_O, const_cast<void*>(gpu_ptr<void>(d_o))},
+      {O, gpu_ptr<void>(o)},
+      {STATS, gpu_ptr<void>(stats)},
+      {D_O, gpu_ptr<void>(d_o)},
       {D_Q, gpu_ptr<void>(d_q)},
       {D_K, gpu_ptr<void>(d_k)},
       {D_V, gpu_ptr<void>(d_v)}};
+  if (mask_arr) {
+    variant_pack[BIAS] = gpu_ptr<void>(*mask_arr);
+  }
 
-  execute_graph(encoder, handle, graph, variant_pack);
+  CHECK_CUDNN_FE_ERROR(graph.encode_graph(encoder, std::move(variant_pack)));
 }
 
 // Defined in scaled_dot_product_attention.cu file.
@@ -469,7 +411,11 @@ bool ScaledDotProductAttention::use_fallback(
 
   return !supports_sdpa_vector(
              q, k, v, has_mask, has_arr_mask, do_causal, output_logsumexp) &&
-      !supports_sdpa_cudnn(q, k, v, has_mask, do_causal, s);
+      !supports_sdpa_cudnn(q, k, v, do_causal, s);
+}
+
+bool ScaledDotProductAttention::supports_bool_mask() {
+  return false;
 }
 
 void ScaledDotProductAttention::eval_gpu(
@@ -487,6 +433,11 @@ void ScaledDotProductAttention::eval_gpu(
   bool has_mask = inputs.size() - has_sinks_ > 3;
   bool has_arr_mask = has_mask && !do_causal_;
 
+  std::optional<array> mask_arr;
+  if (has_arr_mask) {
+    mask_arr = prepare_sdpa_input(inputs[3], s);
+  }
+
   if (supports_sdpa_vector(
           q, k, v, has_mask, has_arr_mask, do_causal_, output_logsumexp_)) {
     if (has_sinks_) {
@@ -495,7 +446,17 @@ void ScaledDotProductAttention::eval_gpu(
       sdpa_vector(q, k, v, scale_, out, do_causal_, std::nullopt, s);
     }
   } else {
-    sdpa_cudnn(q, k, v, scale_, out, stats, do_causal_, output_logsumexp_, s);
+    sdpa_cudnn(
+        q,
+        k,
+        v,
+        scale_,
+        out,
+        stats,
+        do_causal_,
+        mask_arr,
+        output_logsumexp_,
+        s);
   }
 }
 
@@ -515,13 +476,21 @@ void ScaledDotProductAttentionVJP::eval_gpu(
 
   auto& s = stream();
 
-  assert(inputs.size() == 6);
+  assert(inputs.size() >= 6);
+  int primals_size = inputs.size() - 3;
+  bool has_arr_mask = primals_size > 3 + has_sinks_;
+
   array q = prepare_sdpa_input(inputs[0], s);
   array k = prepare_sdpa_input(inputs[1], s);
   array v = prepare_sdpa_input(inputs[2], s);
-  array o = prepare_sdpa_input(inputs[3], s);
-  array stats = prepare_sdpa_input(inputs[4], s);
-  array d_o = prepare_sdpa_input(inputs[5], s);
+  array o = prepare_sdpa_input(inputs[primals_size], s);
+  array stats = prepare_sdpa_input(inputs[primals_size + 1], s);
+  array d_o = prepare_sdpa_input(inputs[primals_size + 2], s);
+
+  std::optional<array> mask_arr;
+  if (has_arr_mask) {
+    mask_arr = prepare_sdpa_input(inputs[3], s);
+  }
 
   assert(outputs.size() == 3);
   auto& d_q = outputs[0];
@@ -529,7 +498,7 @@ void ScaledDotProductAttentionVJP::eval_gpu(
   auto& d_v = outputs[2];
 
   sdpa_backward_cudnn(
-      q, k, v, scale_, o, stats, do_causal_, d_o, d_q, d_k, d_v, s);
+      q, k, v, scale_, o, stats, do_causal_, mask_arr, d_o, d_q, d_k, d_v, s);
 }
 
 } // namespace fast

mlx/backend/cuda/steel/tiles.cuh

@@ -3,31 +3,10 @@
 #pragma once
 
 #include "mlx/backend/cuda/steel/utils.cuh"
+#include "mlx/backend/cuda/vector_types.cuh"
 
 namespace mlx::core::cu {
 
-// Map types to their vector of 2 type float -> float2, double -> double2 etc
-template <typename T>
-struct Vector2;
-template <>
-struct Vector2<double> {
-  using type = double2;
-};
-template <>
-struct Vector2<float> {
-  using type = float2;
-};
-template <>
-struct Vector2<__half> {
-  using type = __half2;
-};
-template <>
-struct Vector2<__nv_bfloat16> {
-  using type = __nv_bfloat162;
-};
-template <typename T>
-using Vector2_t = typename Vector2<T>::type;
-
 /**
  * The basic building block for Ampere mmas. A 16x16 tile distributed across
  * the warp.

mlx/backend/cuda/utils.cpp

@@ -5,6 +5,7 @@
 #include "mlx/dtype_utils.h"
 
 #include <fmt/format.h>
+#include <vector>
 
 namespace mlx::core {
 
@@ -31,6 +32,13 @@ void check_cuda_error(const char* name, CUresult err) {
   }
 }
 
+void check_cudnn_error(const char* name, cudnnStatus_t err) {
+  if (err != CUDNN_STATUS_SUCCESS) {
+    throw std::runtime_error(
+        fmt::format("{} failed: {}.", name, cudnnGetErrorString(err)));
+  }
+}
+
 const char* dtype_to_cuda_type(const Dtype& dtype) {
   switch (dtype) {
     case bool_:
@@ -72,7 +80,6 @@ CudaGraph::CudaGraph(cu::Device& device) {
 }
 
 void CudaGraph::end_capture(cudaStream_t stream) {
-  assert(handle_ == nullptr);
   CHECK_CUDA_ERROR(cudaStreamEndCapture(stream, &handle_));
 }
 
@@ -86,4 +93,23 @@ CudaStream::CudaStream(cu::Device& device) {
   CHECK_CUDA_ERROR(cudaStreamCreateWithFlags(&handle_, cudaStreamNonBlocking));
 }
 
+void* allocate_workspace(cu::CommandEncoder& encoder, size_t workspace_size) {
+  if (workspace_size == 0) {
+    return nullptr;
+  }
+
+  // Workspace allocation should not be captured.
+#ifndef NDEBUG
+  cudaStreamCaptureStatus status;
+  CHECK_CUDA_ERROR(cudaStreamIsCapturing(encoder.stream(), &status));
+  assert(status == cudaStreamCaptureStatusNone);
+#endif
+
+  // Ensure workspace is 256-byte aligned.
+  int nbytes = cuda::ceil_div(workspace_size, 256) * 256;
+  array workspace(cu::malloc_async(nbytes, encoder), {nbytes}, int8);
+  encoder.add_temporary(workspace);
+  return gpu_ptr<void>(workspace);
+}
+
 } // namespace mlx::core

mlx/backend/cuda/utils.h

@@ -31,8 +31,10 @@ inline T* gpu_ptr(array& arr) {
       arr.offset());
 }
 
+// For const array, keep constness in pointer unless it is untyped.
 template <typename T>
-inline const T* gpu_ptr(const array& arr) {
+inline std::conditional_t<std::is_same_v<T, void>, void*, const T*> gpu_ptr(
+    const array& arr) {
   return gpu_ptr<T>(const_cast<array&>(arr));
 }
 
@@ -41,4 +43,7 @@ struct Dtype;
 // Convert Dtype to CUDA C++ types.
 const char* dtype_to_cuda_type(const Dtype& dtype);
 
+// Allocate an empty array and add it as temporary.
+void* allocate_workspace(cu::CommandEncoder& encoder, size_t workspace_size);
+
 } // namespace mlx::core

mlx/backend/cuda/vector_types.cuh

@@ -0,0 +1,48 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+namespace mlx::core::cu {
+
+template <typename T>
+struct Vector2;
+
+template <>
+struct Vector2<double> {
+  using type = double2;
+};
+
+template <>
+struct Vector2<float> {
+  using type = float2;
+};
+
+template <>
+struct Vector2<__half> {
+  using type = __half2;
+};
+
+template <>
+struct Vector2<__nv_bfloat16> {
+  using type = __nv_bfloat162;
+};
+
+template <typename T>
+using Vector2_t = typename Vector2<T>::type;
+
+template <typename T>
+struct Vector4 {
+  T x, y, z, w;
+};
+
+template <typename T>
+using Vector4_t = Vector4<T>;
+
+using bf16x4 = Vector4_t<__nv_bfloat16>;
+using fp16x4 = Vector4_t<__half>;
+using fp32x4 = Vector4_t<float>;
+
+} // namespace mlx::core::cu

mlx/backend/gpu/copy.cpp

@@ -7,8 +7,6 @@
 
 namespace mlx::core {
 
-void copy_gpu(const array& in, array& out, CopyType ctype, const Stream& s);
-
 void copy_gpu(const array& in, array& out, CopyType ctype) {
   copy_gpu(in, out, ctype, out.primitive().stream());
 }

mlx/backend/metal/CMakeLists.txt

@@ -22,7 +22,7 @@ function(make_jit_source SRC_FILE)
 endfunction(make_jit_source)
 
 make_jit_source(utils kernels/bf16.h kernels/bf16_math.h kernels/complex.h
-                kernels/defines.h)
+                kernels/defines.h kernels/logging.h)
 make_jit_source(unary_ops kernels/erf.h kernels/expm1f.h kernels/fp8.h)
 make_jit_source(binary_ops)
 make_jit_source(ternary_ops)
@@ -79,6 +79,21 @@ if(MLX_METAL_JIT)
   make_jit_source(fp_quantized kernels/quantized_utils.h kernels/fp8.h
                   kernels/fp4.h)
   make_jit_source(gemv_masked)
+
+  make_jit_source(steel/attn/kernels/steel_attention)
+
+  make_jit_source(
+    steel/gemm/gemm_nax kernels/steel/utils.h kernels/steel/gemm/nax.h
+    kernels/steel/gemm/params.h kernels/steel/gemm/transforms.h)
+  make_jit_source(steel/gemm/kernels/steel_gemm_fused_nax)
+  make_jit_source(steel/gemm/kernels/steel_gemm_gather_nax)
+
+  make_jit_source(quantized_nax kernels/quantized_utils.h)
+  make_jit_source(fp_quantized_nax kernels/quantized_utils.h kernels/fp8.h
+                  kernels/fp4.h)
+
+  make_jit_source(steel/attn/kernels/steel_attention_nax)
+
 else()
   target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/nojit_kernels.cpp)
 endif()
@@ -121,14 +136,6 @@ if(NOT MLX_METAL_PATH)
   set(MLX_METAL_PATH ${CMAKE_CURRENT_BINARY_DIR}/kernels/)
 endif()
 
-if((MLX_METAL_VERSION GREATER_EQUAL 400) AND (MACOS_SDK_VERSION GREATER_EQUAL
-                                              26.2))
-  set(MLX_ENABLE_NAX TRUE)
-  target_compile_definitions(mlx PRIVATE MLX_ENABLE_NAX)
-else()
-  set(MLX_ENABLE_NAX FALSE)
-endif()
-
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/kernels)
 
 target_compile_definitions(mlx

mlx/backend/metal/allocator.cpp

@@ -149,7 +149,9 @@ Buffer MetalAllocator::malloc(size_t size) {
       buf = device_->newBuffer(size, resource_options);
     }
     if (!buf) {
-      return Buffer{nullptr};
+      std::ostringstream msg;
+      msg << "[malloc] Unable to allocate " << size << " bytes.";
+      throw std::runtime_error(msg.str());
     }
     lk.lock();
     num_resources_++;
@@ -201,6 +203,32 @@ size_t MetalAllocator::size(Buffer buffer) const {
   return static_cast<MTL::Buffer*>(buffer.ptr())->length();
 }
 
+Buffer MetalAllocator::make_buffer(void* ptr, size_t size) {
+  auto buf = device_->newBuffer(ptr, size, resource_options, nullptr);
+  if (!buf) {
+    return Buffer{nullptr};
+  }
+  std::unique_lock lk(mutex_);
+  residency_set_.insert(buf);
+  active_memory_ += buf->length();
+  peak_memory_ = std::max(peak_memory_, active_memory_);
+  num_resources_++;
+  return Buffer{static_cast<void*>(buf)};
+}
+
+void MetalAllocator::release(Buffer buffer) {
+  auto buf = static_cast<MTL::Buffer*>(buffer.ptr());
+  if (buf == nullptr) {
+    return;
+  }
+  std::unique_lock lk(mutex_);
+  active_memory_ -= buf->length();
+  num_resources_--;
+  lk.unlock();
+  auto pool = metal::new_scoped_memory_pool();
+  buf->release();
+}
+
 MetalAllocator& allocator() {
   // By creating the |allocator_| on heap, the destructor of MetalAllocator
   // will not be called on exit and buffers in the cache will be leaked. This

mlx/backend/metal/allocator.h

@@ -21,6 +21,9 @@ class MetalAllocator : public allocator::Allocator {
   virtual Buffer malloc(size_t size) override;
   virtual void free(Buffer buffer) override;
   virtual size_t size(Buffer buffer) const override;
+  virtual Buffer make_buffer(void* ptr, size_t size) override;
+  virtual void release(Buffer buffer) override;
+
   size_t get_active_memory() {
     return active_memory_;
   };

mlx/backend/metal/device.cpp

@@ -224,7 +224,7 @@ MTL::Library* load_library(
   std::ostringstream msg;
   msg << "Failed to load the metallib " << lib_name << ".metallib. "
       << "We attempted to load it from <" << current_binary_dir() << "/"
-      << lib_name << ".metallib" << ">";
+      << lib_name << ".metallib>";
 #ifdef SWIFTPM_BUNDLE
   msg << " and from the Swift PM bundle.";
 #endif
@@ -529,6 +529,11 @@ MTL::Library* Device::build_library_(const std::string& source_string) {
   auto options = MTL::CompileOptions::alloc()->init();
   options->setFastMathEnabled(false);
   options->setLanguageVersion(get_metal_version());
+#ifndef NDEBUG
+  if (options->languageVersion() >= MTL::LanguageVersion3_2) {
+    options->setEnableLogging(true);
+  }
+#endif
   auto mtl_lib = device_->newLibrary(ns_code, options, &error);
   options->release();
 

mlx/backend/metal/device.h

@@ -265,14 +265,19 @@ Device& device(mlx::core::Device);
 
 std::unique_ptr<void, std::function<void(void*)>> new_scoped_memory_pool();
 
-#ifdef MLX_ENABLE_NAX
-
 inline bool is_nax_available() {
-  static bool is_nax_available_ =
-      metal::device(mlx::core::Device::gpu).get_architecture_gen() >= 17;
+  auto _check_nax = []() {
+    bool can_use_nax = false;
+    if (__builtin_available(
+            macOS 26.2, iOS 26.2, tvOS 26.2, visionOS 26.2, *)) {
+      can_use_nax = true;
+    }
+    can_use_nax &=
+        metal::device(mlx::core::Device::gpu).get_architecture_gen() >= 17;
+    return can_use_nax;
+  };
+  static bool is_nax_available_ = _check_nax();
   return is_nax_available_;
 }
 
-#endif // MLX_ENABLE_NAX
-
 } // namespace mlx::core::metal

mlx/backend/metal/jit/includes.h

@@ -43,5 +43,15 @@ const char* conv();
 const char* steel_conv();
 const char* steel_conv_general();
 const char* gemv_masked();
+const char* steel_attention();
+
+const char* gemm_nax();
+const char* steel_gemm_fused_nax();
+const char* steel_gemm_gather_nax();
+
+const char* quantized_nax();
+const char* fp_quantized_nax();
+
+const char* steel_attention_nax();
 
 } // namespace mlx::core::metal

mlx/backend/metal/jit_kernels.cpp

@@ -877,4 +877,208 @@ MTL::ComputePipelineState* get_gather_qmm_kernel(
   return d.get_kernel(kernel_name, lib, hash_name, func_consts);
 }
 
+MTL::ComputePipelineState* get_steel_gemm_fused_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::ostringstream kernel_source;
+    kernel_source << metal::utils() << metal::gemm_nax()
+                  << metal::steel_gemm_fused_nax()
+                  << get_template_definition(
+                         lib_name,
+                         "gemm",
+                         get_type_string(out.dtype()),
+                         bm,
+                         bn,
+                         bk,
+                         wm,
+                         wn,
+                         transpose_a,
+                         transpose_b);
+    return kernel_source.str();
+  });
+  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_steel_gemm_gather_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool rhs) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source;
+    concatenate(
+        kernel_source,
+        metal::utils(),
+        metal::gemm_nax(),
+        metal::steel_gemm_gather_nax(),
+        get_template_definition(
+            lib_name,
+            rhs ? "gather_mm_rhs_nax" : "gather_mm_nax",
+            get_type_string(out.dtype()),
+            bm,
+            bn,
+            bk,
+            wm,
+            wn,
+            transpose_a,
+            transpose_b));
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_qmm_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& template_def,
+    const std::string& mode) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source;
+    concatenate(
+        kernel_source,
+        metal::utils(),
+        metal::gemm_nax(),
+        metal::quantized_utils(),
+        (mode == "affine") ? metal::quantized_nax() : metal::fp_quantized_nax(),
+        template_def);
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib);
+}
+
+MTL::ComputePipelineState* get_gather_qmm_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& x,
+    int group_size,
+    int bits,
+    const std::string& mode,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool transpose) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source;
+    concatenate(
+        kernel_source,
+        metal::utils(),
+        metal::gemm_nax(),
+        metal::quantized_utils());
+    bool is_affine = mode == "affine";
+    concatenate(
+        kernel_source,
+        is_affine ? metal::quantized_nax() : metal::fp_quantized_nax(),
+        get_template_definition(
+            lib_name,
+            (is_affine ? "affine" : "fp") + std::string("_gather_qmm_rhs_nax"),
+            get_type_string(x.dtype()),
+            group_size,
+            bits,
+            bm,
+            bn,
+            bk,
+            wm,
+            wn,
+            transpose));
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_steel_attention_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& q,
+    int bq,
+    int bk,
+    int bd,
+    int wm,
+    int wn,
+    const array& m) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source;
+    concatenate(
+        kernel_source,
+        metal::utils(),
+        metal::steel_attention(),
+        get_template_definition(
+            lib_name,
+            "attention",
+            get_type_string(q.dtype()),
+            bq,
+            bk,
+            bd,
+            wm,
+            wn,
+            get_type_string(m.dtype())));
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
+}
+
+MTL::ComputePipelineState* get_steel_attention_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& q,
+    int bq,
+    int bk,
+    int bd,
+    int wm,
+    int wn,
+    const array& m) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::string kernel_source;
+    concatenate(
+        kernel_source,
+        metal::utils(),
+        metal::steel_attention_nax(),
+        get_template_definition(
+            lib_name,
+            "attention_nax",
+            get_type_string(q.dtype()),
+            bq,
+            bk,
+            bd,
+            wm,
+            wn,
+            get_type_string(m.dtype())));
+    return kernel_source;
+  });
+  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
+}
+
 } // namespace mlx::core

mlx/backend/metal/kernels.h

@@ -257,6 +257,83 @@ MTL::ComputePipelineState* get_gather_qmm_kernel(
     int wn,
     bool transpose);
 
+MTL::ComputePipelineState* get_steel_gemm_fused_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn);
+
+MTL::ComputePipelineState* get_steel_gemm_gather_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool rhs);
+
+MTL::ComputePipelineState* get_qmm_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& template_def,
+    const std::string& mode);
+
+MTL::ComputePipelineState* get_gather_qmm_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& x,
+    int group_size,
+    int bits,
+    const std::string& mode,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn,
+    bool transpose);
+
+MTL::ComputePipelineState* get_steel_attention_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& q,
+    int bq,
+    int bk,
+    int bd,
+    int wm,
+    int wn,
+    const array& m);
+
+MTL::ComputePipelineState* get_steel_attention_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& q,
+    int bq,
+    int bk,
+    int bd,
+    int wm,
+    int wn,
+    const array& m);
+
 // Create a GPU kernel template definition for JIT compilation
 template <typename... Args>
 std::string get_template_definition(

mlx/backend/metal/kernels/CMakeLists.txt

@@ -6,15 +6,23 @@ set(BASE_HEADERS
     erf.h
     expm1f.h
     fp8.h
+    logging.h
     utils.h)
 
 function(build_kernel_base TARGET SRCFILE DEPS)
-  set(METAL_FLAGS -x metal -Wall -Wextra -fno-fast-math -Wno-c++17-extensions)
+  set(METAL_FLAGS
+      -x
+      metal
+      -Wall
+      -Wextra
+      -fno-fast-math
+      -Wno-c++17-extensions
+      -Wno-c++20-extensions)
   if(MLX_METAL_DEBUG)
     set(METAL_FLAGS ${METAL_FLAGS} -gline-tables-only -frecord-sources)
   endif()
-  if(MLX_ENABLE_NAX)
-    set(METAL_FLAGS ${METAL_FLAGS} -Wno-c++20-extensions -std=metal4.0)
+  if(CMAKE_BUILD_TYPE STREQUAL "Debug" AND MLX_METAL_VERSION GREATER_EQUAL 320)
+    set(METAL_FLAGS ${METAL_FLAGS} -fmetal-enable-logging)
   endif()
   if(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "")
     set(METAL_FLAGS ${METAL_FLAGS}
@@ -63,6 +71,7 @@ set(STEEL_HEADERS
     steel/gemm/gemm.h
     steel/gemm/mma.h
     steel/gemm/loader.h
+    steel/gemm/params.h
     steel/gemm/transforms.h
     steel/gemm/kernels/steel_gemm_fused.h
     steel/gemm/kernels/steel_gemm_gather.h
@@ -88,7 +97,26 @@ set(STEEL_ATTN_HEADERS
     steel/attn/transforms.h
     steel/attn/kernels/steel_attention.h)
 
-build_kernel(steel/attn/kernels/steel_attention ${STEEL_ATTN_HEADERS})
+set(STEEL_NAX_HEADERS
+    steel/defines.h
+    steel/utils.h
+    steel/gemm/params.h
+    steel/gemm/transforms.h
+    steel/gemm/nax.h
+    steel/gemm/gemm_nax.h
+    steel/utils/type_traits.h
+    steel/utils/integral_constant.h
+    steel/gemm/kernels/steel_gemm_fused_nax.h
+    steel/gemm/kernels/steel_gemm_gather_nax.h)
+
+set(STEEL_NAX_ATTN_HEADERS
+    steel/defines.h
+    steel/utils.h
+    steel/attn/nax.h
+    steel/utils/type_traits.h
+    steel/utils/integral_constant.h
+    steel/attn/params.h
+    steel/attn/kernels/steel_attention_nax.h)
 
 if(NOT MLX_METAL_JIT)
   build_kernel(arange arange.h)
@@ -105,7 +133,7 @@ if(NOT MLX_METAL_JIT)
     reduction/reduce_col.h
     reduction/reduce_row.h)
   build_kernel(quantized quantized.h quantized_utils.h ${STEEL_HEADERS})
-  build_kernel(fp_quantized fp4.h fp_quantized.h quantized_utils.h
+  build_kernel(fp_quantized fp4.h fp8.h fp_quantized.h quantized_utils.h
                ${STEEL_HEADERS})
   build_kernel(scan scan.h)
   build_kernel(softmax softmax.h)
@@ -121,30 +149,23 @@ if(NOT MLX_METAL_JIT)
   build_kernel(steel/gemm/kernels/steel_gemm_splitk ${STEEL_HEADERS})
   build_kernel(steel/gemm/kernels/steel_gemm_segmented ${STEEL_HEADERS})
   build_kernel(gemv_masked steel/utils.h)
-endif()
+  build_kernel(steel/attn/kernels/steel_attention ${STEEL_ATTN_HEADERS})
 
-if(MLX_ENABLE_NAX)
+  if((MLX_METAL_VERSION GREATER_EQUAL 400) AND (MACOS_SDK_VERSION GREATER_EQUAL
+                                                26.2))
 
-  set(STEEL_NAX_HEADERS
-      steel/defines.h
-      steel/utils.h
-      steel/gemm/transforms.h
-      steel/gemm/nax.h
-      steel/gemm/gemm_nax.h
-      steel/utils/type_traits.h
-      steel/utils/integral_constant.h)
+    build_kernel(steel/gemm/kernels/steel_gemm_fused_nax ${STEEL_NAX_HEADERS})
+    build_kernel(steel/gemm/kernels/steel_gemm_gather_nax ${STEEL_NAX_HEADERS})
 
-  build_kernel(steel/gemm/kernels/steel_gemm_fused_nax ${STEEL_NAX_HEADERS})
-  build_kernel(steel/gemm/kernels/steel_gemm_gather_nax ${STEEL_NAX_HEADERS})
+    build_kernel(quantized_nax quantized_nax.h ${STEEL_NAX_HEADERS})
+    build_kernel(fp_quantized_nax fp4.h fp8.h fp_quantized_nax.h
+                 ${STEEL_NAX_HEADERS})
 
-  build_kernel(quantized_nax quantized_nax.h ${STEEL_NAX_HEADERS})
-  build_kernel(fp_quantized_nax fp_quantized_nax.h ${STEEL_NAX_HEADERS})
+    build_kernel(steel/attn/kernels/steel_attention_nax
+                 ${STEEL_NAX_ATTN_HEADERS})
 
-  set(STEEL_NAX_ATTN_HEADERS
-      steel/defines.h steel/utils.h steel/attn/nax.h steel/utils/type_traits.h
-      steel/utils/integral_constant.h)
+  endif()
 
-  build_kernel(steel/attn/kernels/steel_attention_nax ${STEEL_NAX_ATTN_HEADERS})
 endif()
 
 add_custom_command(

mlx/backend/metal/kernels/binary_ops.h

@@ -5,6 +5,8 @@
 #include <metal_integer>
 #include <metal_math>
 
+constant mlx::os_log logger("mlx", "binary_ops");
+
 struct Add {
   template <typename T>
   T operator()(T x, T y) {
@@ -225,6 +227,8 @@ struct Power {
     T res = 1;
     // Undefined to raise integer to negative power
     if (exp < 0) {
+      logger.log_debug(
+          "int pow exp<0 (base=%ld exp=%ld)", (long)base, (long)exp);
       return 0;
     }
 

mlx/backend/metal/kernels/fp4.h

@@ -1,23 +1,5 @@
 #pragma once
 
-constexpr constant static float FP4_LUT[16] = {
-    +0.0f,
-    +0.5f,
-    +1.0f,
-    +1.5f,
-    +2.0f,
-    +3.0f,
-    +4.0f,
-    +6.0f,
-    -0.0f,
-    -0.5f,
-    -1.0f,
-    -1.5f,
-    -2.0f,
-    -3.0f,
-    -4.0f,
-    -6.0f};
-
 struct fp4_e2m1 {
   fp4_e2m1(float x) {
     if (metal::isnan(x)) {
@@ -48,11 +30,18 @@ struct fp4_e2m1 {
     bits |= sign_bit;
   }
 
-  operator float() {
+  operator float16_t() {
     half converted = as_type<half>(ushort((bits & 7) << 9));
     converted *= 16384.0;
-    converted = bits & 8 ? -converted : converted;
-    return converted;
+    return bits & 8 ? -converted : converted;
+  }
+
+  operator float() {
+    return static_cast<float>(this->operator float16_t());
+  }
+
+  operator bfloat16_t() {
+    return static_cast<bfloat16_t>(this->operator float16_t());
   }
 
   uint8_t bits;

mlx/backend/metal/kernels/fp8.h

@@ -29,24 +29,20 @@ struct fp8_e4m3 {
     bits |= static_cast<uint8_t>(sign >> 24);
   }
 
+  operator float16_t() {
+    uint16_t v = (bits & 127) << 7;
+    half converted = as_type<half>(v);
+    converted *= 256.0;
+    auto sign = bits & 128;
+    return (sign ? -converted : converted);
+  }
+
+  operator bfloat16_t() {
+    return static_cast<bfloat16_t>(this->operator float16_t());
+  }
+
   operator float() {
-    // From PyTorch:
-    // https://github.com/pytorch/pytorch/blob/e3643e1e0e923f0fc063dfab6f45c956d568919d/c10/util/Float8_e4m3fn.h#L46
-    uint32_t w = static_cast<uint32_t>(bits) << 24;
-    uint32_t sign = w & 0x80000000;
-    uint32_t nonsign = w & 0x7FFFFFFF;
-
-    uint32_t renorm_shift = metal::clz(nonsign);
-    renorm_shift = renorm_shift > 4 ? renorm_shift - 4 : 0;
-
-    int32_t inf_nan_mask =
-        (static_cast<int32_t>(nonsign + 0x01000000) >> 8) & 0x7F800000;
-    int32_t zero_mask = static_cast<int32_t>(nonsign - 1) >> 31;
-    uint32_t result = sign |
-        ((((nonsign << renorm_shift >> 4) + ((0x78 - renorm_shift) << 23)) |
-          inf_nan_mask) &
-         ~zero_mask);
-    return as_type<float>(result);
+    return static_cast<float>(this->operator float16_t());
   }
 
   uint8_t bits;
@@ -74,8 +70,10 @@ struct fp8_e8m0 {
     uint16_t out = (bits == 0 ? 0x40 : (static_cast<uint16_t>(bits) << 7));
     return as_type<bfloat16_t>(out);
   }
+
   operator float() {
-    return static_cast<float>(this->operator bfloat16_t());
+    uint32_t out = (bits == 0 ? 0x400000 : (static_cast<uint16_t>(bits) << 23));
+    return as_type<float>(out);
   }
 
   uint8_t bits;

mlx/backend/metal/kernels/fp_quantized.h

@@ -17,9 +17,9 @@ using namespace metal;
 MLX_MTL_CONST int SIMD_SIZE = 32;
 MLX_MTL_CONST int QUAD_SIZE = 4;
 
-template <int wsize = 8>
+template <int wsize = 8, int bits = 4>
 inline constexpr short get_pack_factor() {
-  return wsize / 4;
+  return wsize / bits;
 }
 
 template <int wsize = 8>
@@ -27,9 +27,14 @@ inline constexpr short get_bytes_per_pack() {
   return wsize / 8;
 }
 
-template <typename T>
+template <typename T, int group_size>
 static inline T dequantize_scale(uint8_t s) {
-  return T(*(thread fp8_e8m0*)(&s));
+  if constexpr (group_size == 16) {
+    // Use nv scale
+    return T(*(thread fp8_e4m3*)(&s));
+  } else {
+    return T(*(thread fp8_e8m0*)(&s));
+  }
 }
 
 template <int bits>
@@ -43,34 +48,29 @@ struct Quantize {
   }
 };
 
-template <int bits>
+template <int bits, typename U = float>
 struct Dequantize {
-  float operator()(uint8_t x) {
-    if (bits == 8) {
-      return float(*(thread fp8_e4m3*)(&x));
+  U operator()(uint8_t x) {
+    if constexpr (bits == 8) {
+      return U(*(thread fp8_e4m3*)(&x));
     } else {
-      return float(*(thread fp4_e2m1*)(&x));
+      return U(*(thread fp4_e2m1*)(&x));
     }
   }
 };
 
 template <typename T, typename U, int values_per_thread>
 inline void load_vector(const device T* x, thread U* x_thread) {
-  for (int i = 0; i < values_per_thread; i += 4) {
+#pragma unroll
+  for (int i = 0; i < values_per_thread; i++) {
     x_thread[i] = x[i];
-    x_thread[i + 1] = x[i + 1];
-    x_thread[i + 2] = x[i + 2];
-    x_thread[i + 3] = x[i + 3];
   }
 }
 
 template <typename T, typename U, int values_per_thread>
 inline void load_vector_safe(const device T* x, thread U* x_thread, int N) {
-  for (int i = 0; i < N; i += 4) {
+  for (int i = 0; i < N; i++) {
     x_thread[i] = x[i];
-    x_thread[i + 1] = x[i + 1];
-    x_thread[i + 2] = x[i + 2];
-    x_thread[i + 3] = x[i + 3];
   }
 
   for (int i = N; i < values_per_thread; i++) {
@@ -78,53 +78,70 @@ inline void load_vector_safe(const device T* x, thread U* x_thread, int N) {
   }
 }
 
-template <typename U, int values_per_thread>
+template <typename U, int values_per_thread, int bits>
 inline U qdot(const device uint8_t* w, const thread U* x_thread, U scale) {
   U accum = 0;
-  const device uint16_t* ws = (const device uint16_t*)w;
-  for (int i = 0; i < (values_per_thread / 4); i++) {
-    accum +=
-        (x_thread[4 * i] * Dequantize<4>{}(ws[i]) +
-         x_thread[4 * i + 1] * Dequantize<4>{}(ws[i] >> 4) +
-         x_thread[4 * i + 2] * Dequantize<4>{}(ws[i] >> 8) +
-         x_thread[4 * i + 3] * Dequantize<4>{}(ws[i] >> 12));
+  if constexpr (bits == 4) {
+    const device uint16_t* ws = (const device uint16_t*)w;
+    for (int i = 0; i < (values_per_thread / 4); i++) {
+      accum +=
+          (x_thread[4 * i] * Dequantize<4>{}(ws[i]) +
+           x_thread[4 * i + 1] * Dequantize<4>{}(ws[i] >> 4) +
+           x_thread[4 * i + 2] * Dequantize<4>{}(ws[i] >> 8) +
+           x_thread[4 * i + 3] * Dequantize<4>{}(ws[i] >> 12));
+    }
+  } else {
+    for (int i = 0; i < values_per_thread; i++) {
+      accum += x_thread[i] * Dequantize<8>{}(w[i]);
+    }
   }
+
   return scale * accum;
 }
 
-template <typename U, int values_per_thread>
+template <typename U, int values_per_thread, int bits>
 inline U
 qdot_safe(const device uint8_t* w, const thread U* x_thread, U scale, int N) {
   U accum = 0;
 
-  const device uint16_t* ws = (const device uint16_t*)w;
-  for (int i = 0; i < (N / 4); i++) {
-    accum +=
-        (x_thread[4 * i] * Dequantize<4>{}(ws[i]) +
-         x_thread[4 * i + 1] * Dequantize<4>{}(ws[i] >> 4) +
-         x_thread[4 * i + 2] * Dequantize<4>{}(ws[i] >> 8) +
-         x_thread[4 * i + 3] * Dequantize<4>{}(ws[i] >> 12));
+  if constexpr (bits == 4) {
+    const device uint16_t* ws = (const device uint16_t*)w;
+    for (int i = 0; i < (N / 4); i++) {
+      accum +=
+          (x_thread[4 * i] * Dequantize<4>{}(ws[i]) +
+           x_thread[4 * i + 1] * Dequantize<4>{}(ws[i] >> 4) +
+           x_thread[4 * i + 2] * Dequantize<4>{}(ws[i] >> 8) +
+           x_thread[4 * i + 3] * Dequantize<4>{}(ws[i] >> 12));
+    }
+  } else {
+    for (int i = 0; i < N; i++) {
+      accum += x_thread[i] * Dequantize<8>{}(w[i]);
+    }
   }
   return scale * accum;
 }
 
-template <typename U, int values_per_thread>
+template <typename U, int values_per_thread, int bits>
 inline void qouter(const thread uint8_t* w, U x, U scale, thread U* result) {
-  for (int i = 0; i < (values_per_thread / 2); i++) {
-    result[2 * i] += x * scale * Dequantize<4>{}(w[i]);
-    result[2 * i + 1] += x * scale * Dequantize<4>{}(w[i] >> 4);
+  if constexpr (bits == 4) {
+    for (int i = 0; i < (values_per_thread / 2); i++) {
+      result[2 * i] += x * scale * Dequantize<4>{}(w[i]);
+      result[2 * i + 1] += x * scale * Dequantize<4>{}(w[i] >> 4);
+    }
+  } else {
+    for (int i = 0; i < values_per_thread; i++) {
+      result[i] += x * scale * Dequantize<8>{}(w[i]);
+    }
   }
 }
 
-template <typename U, int N>
-inline void dequantize(
-    const device uint8_t* w,
-    U scale,
-    threadgroup U* w_local,
-    const threadgroup U* lut) {
-  for (int i = 0; i < (N / 2); i++) {
-    w_local[2 * i] = scale * lut[w[i] & 0xf];
-    w_local[2 * i + 1] = scale * lut[(w[i] >> 4) & 0xf];
+template <typename U, int bits>
+inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) {
+  if constexpr (bits == 4) {
+    w_local[0] = scale * Dequantize<4, U>{}(w);
+    w_local[1] = scale * Dequantize<4, U>{}(w >> 4);
+  } else {
+    w_local[0] = scale * Dequantize<8, U>{}(w);
   }
 }
 
@@ -135,21 +152,20 @@ template <
     short dst_ld,
     short reduction_dim,
     short tgp_size,
-    short group_size>
+    short group_size,
+    short bits>
 struct QuantizedBlockLoader {
-  static_assert(
-      BCOLS <= group_size,
-      "The group size should be larger than the columns");
-  static_assert(
-      group_size % BCOLS == 0,
-      "The group size should be divisible by the columns");
-
-  MLX_MTL_CONST short pack_factor = get_pack_factor<8>();
+  MLX_MTL_CONST short pack_factor = get_pack_factor<8, bits>();
   MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack();
   MLX_MTL_CONST short BCOLS_PACKED = BCOLS / pack_factor;
   MLX_MTL_CONST short n_reads =
       (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size;
-  MLX_MTL_CONST short group_steps = group_size / BCOLS;
+  MLX_MTL_CONST short group_steps = group_size < BCOLS ? 1 : group_size / BCOLS;
+  MLX_MTL_CONST short scale_step = group_size < BCOLS ? BCOLS / group_size : 1;
+
+  static_assert(
+      (n_reads * pack_factor) <= group_size,
+      "The number of reads per thread must be less than the group size.");
 
   const int src_ld;
   const int tile_stride;
@@ -163,14 +179,12 @@ struct QuantizedBlockLoader {
   threadgroup T* dst;
   const device uint8_t* src;
   const device uint8_t* scales;
-  threadgroup T* lut;
 
   QuantizedBlockLoader(
       const device uint8_t* src_,
       const device uint8_t* scales_,
       const int src_ld_,
       threadgroup T* dst_,
-      threadgroup T* lut_,
       ushort simd_group_id [[simdgroup_index_in_threadgroup]],
       ushort simd_lane_id [[thread_index_in_simdgroup]])
       : src_ld(src_ld_),
@@ -185,23 +199,19 @@ struct QuantizedBlockLoader {
         dst(dst_ + bi * dst_ld + bj * pack_factor),
         src(src_ + bi * src_ld * bytes_per_pack / pack_factor +
             bj * bytes_per_pack),
-        scales(scales_ + bi * src_ld / group_size),
-        lut(lut_) {
-    if (simd_group_id == 0 && simd_lane_id < 16) {
-      lut[simd_lane_id] = static_cast<T>(FP4_LUT[simd_lane_id]);
-    }
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-  }
+        scales(
+            scales_ + bi * src_ld / group_size +
+            (bj * pack_factor) / group_size) {}
 
   void load_unsafe() const {
     if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) {
       return;
     }
 
-    T scale = dequantize_scale<T>(*scales);
+    T scale = dequantize_scale<T, group_size>(*scales);
     for (int i = 0; i < n_reads; i++) {
-      dequantize<T, pack_factor>(
-          src + i * bytes_per_pack, scale, dst + i * pack_factor, lut);
+      dequantize<T, bits>(
+          src[i * bytes_per_pack], scale, dst + i * pack_factor);
     }
   }
 
@@ -224,13 +234,10 @@ struct QuantizedBlockLoader {
       return;
     }
 
-    T scale = dequantize_scale<T>(*scales);
+    T scale = dequantize_scale<T, group_size>(*scales);
     for (int i = 0; i < n_reads; i++) {
-      dequantize<T, pack_factor>(
-          (device uint8_t*)(src + i * bytes_per_pack),
-          scale,
-          dst + i * pack_factor,
-          lut);
+      dequantize<T, bits>(
+          src[i * bytes_per_pack], scale, dst + i * pack_factor);
     }
   }
 
@@ -244,7 +251,7 @@ struct QuantizedBlockLoader {
           scales++;
         }
       } else {
-        scales++;
+        scales += scale_step;
       }
     } else {
       scales += group_stride;
@@ -264,10 +271,13 @@ METAL_FUNC void fp_qmv_quad_impl(
     uint quad_gid [[quadgroup_index_in_threadgroup]],
     uint quad_lid [[thread_index_in_quadgroup]]) {
   constexpr int quads_per_simd = SIMD_SIZE / QUAD_SIZE;
-  constexpr int pack_factor = 8;
+  constexpr int pack_factor = get_pack_factor<32, bits>();
   constexpr int values_per_thread = D / QUAD_SIZE;
+  constexpr int steps_per_thread =
+      values_per_thread < group_size ? 1 : values_per_thread / group_size;
+  constexpr int values_per_step = values_per_thread / steps_per_thread;
   constexpr int packs_per_thread = values_per_thread / pack_factor;
-  constexpr int scale_step_per_thread = group_size / values_per_thread;
+  constexpr int packs_per_step = values_per_step / pack_factor;
   constexpr int results_per_quadgroup = 8;
 
   typedef float U;
@@ -281,7 +291,8 @@ METAL_FUNC void fp_qmv_quad_impl(
   const int out_row = tid.y * quads_per_simd * results_per_quadgroup + quad_gid;
 
   w += out_row * in_vec_size_w + quad_lid * packs_per_thread;
-  scales += out_row * in_vec_size_g + quad_lid / scale_step_per_thread;
+  scales +=
+      out_row * in_vec_size_g + (quad_lid * values_per_thread) / group_size;
   x += tid.x * in_vec_size + quad_lid * values_per_thread;
   y += tid.x * out_vec_size + out_row;
 
@@ -290,10 +301,15 @@ METAL_FUNC void fp_qmv_quad_impl(
   for (int row = 0; row < results_per_quadgroup; row++) {
     auto wl = (const device uint8_t*)(w + row * in_vec_size_w * quads_per_simd);
     const device uint8_t* sl = scales + row * in_vec_size_g * quads_per_simd;
-
-    U s = dequantize_scale<U>(sl[0]);
-    if (row * quads_per_simd + out_row < out_vec_size) {
-      result[row] += qdot<U, values_per_thread>(wl, x_thread, s);
+#pragma unroll
+    for (int k = 0; k < steps_per_thread; ++k) {
+      U s = dequantize_scale<U, group_size>(sl[0]);
+      if (row * quads_per_simd + out_row < out_vec_size) {
+        result[row] += qdot<U, values_per_step, bits>(
+            wl, x_thread + k * values_per_step, s);
+      }
+      sl++;
+      wl += (sizeof(uint32_t) / sizeof(uint8_t)) * packs_per_step;
     }
   }
 
@@ -319,7 +335,7 @@ METAL_FUNC void fp_qmv_fast_impl(
   constexpr int packs_per_thread = 2;
   constexpr int num_simdgroups = 2;
   constexpr int results_per_simdgroup = 4;
-  constexpr int pack_factor = get_pack_factor<32>();
+  constexpr int pack_factor = get_pack_factor<32, bits>();
   constexpr int bytes_per_pack = get_bytes_per_pack<32>();
   constexpr int values_per_thread = pack_factor * packs_per_thread;
   constexpr int block_size = values_per_thread * SIMD_SIZE;
@@ -349,8 +365,8 @@ METAL_FUNC void fp_qmv_fast_impl(
       auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
       const device auto* sl = scales + row * in_vec_size_g;
 
-      U s = dequantize_scale<U>(sl[0]);
-      result[row] += qdot<U, values_per_thread>(wl, x_thread, s);
+      U s = dequantize_scale<U, group_size>(sl[0]);
+      result[row] += qdot<U, values_per_thread, bits>(wl, x_thread, s);
     }
 
     ws += block_size * bytes_per_pack / pack_factor;
@@ -380,7 +396,7 @@ METAL_FUNC void fp_qmv_impl(
   constexpr int num_simdgroups = 2;
   constexpr int results_per_simdgroup = 4;
   constexpr int packs_per_thread = 1;
-  constexpr int pack_factor = get_pack_factor<32>();
+  constexpr int pack_factor = get_pack_factor<32, bits>();
   constexpr int bytes_per_pack = get_bytes_per_pack<32>();
 
   constexpr int values_per_thread = pack_factor * packs_per_thread;
@@ -423,7 +439,7 @@ METAL_FUNC void fp_qmv_impl(
         const device auto* sl = scales + row * in_vec_size_g;
 
         uint8_t s = sl[0];
-        result[row] += qdot<U, values_per_thread>(wl, x_thread, s);
+        result[row] += qdot<U, values_per_thread, bits>(wl, x_thread, s);
       }
 
       ws += block_size * bytes_per_pack / pack_factor;
@@ -441,8 +457,8 @@ METAL_FUNC void fp_qmv_impl(
         auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
         const device auto* sl = scales + row * in_vec_size_g;
 
-        U s = dequantize_scale<U>(sl[0]);
-        result[row] += qdot<U, values_per_thread>(wl, x_thread, s);
+        U s = dequantize_scale<U, group_size>(sl[0]);
+        result[row] += qdot<U, values_per_thread, bits>(wl, x_thread, s);
       }
     }
 
@@ -470,8 +486,8 @@ METAL_FUNC void fp_qmv_impl(
         auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
         const device auto* sl = scales + row * in_vec_size_g;
 
-        U s = dequantize_scale<U>(sl[0]);
-        result[row] += qdot<U, values_per_thread>(wl, x_thread, s);
+        U s = dequantize_scale<U, group_size>(sl[0]);
+        result[row] += qdot<U, values_per_thread, bits>(wl, x_thread, s);
       }
 
       ws += block_size * bytes_per_pack / pack_factor;
@@ -489,9 +505,9 @@ METAL_FUNC void fp_qmv_impl(
         auto wl = (const device uint8_t*)(ws + row * in_vec_size_w);
         const device auto* sl = scales + row * in_vec_size_g;
 
-        U s = dequantize_scale<U>(sl[0]);
+        U s = dequantize_scale<U, group_size>(sl[0]);
         result[row] +=
-            qdot_safe<U, values_per_thread>(wl, x_thread, s, remaining);
+            qdot_safe<U, values_per_thread, bits>(wl, x_thread, s, remaining);
       }
     }
     for (int row = 0; row < results_per_simdgroup; row++) {
@@ -515,10 +531,10 @@ METAL_FUNC void fp_qvm_impl(
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
   constexpr int num_simdgroups = 2;
-  constexpr int pack_factor = get_pack_factor<32>();
+  constexpr int pack_factor = get_pack_factor<32, bits>();
   constexpr int bytes_per_pack = get_bytes_per_pack();
 
-  constexpr int tn = 32 / pack_factor;
+  constexpr int tn = group_size / pack_factor;
   constexpr int block_size = SIMD_SIZE;
 
   using W_T = uint32_t;
@@ -537,6 +553,7 @@ METAL_FUNC void fp_qvm_impl(
   // Adjust positions
   const int out_vec_size_w = out_vec_size * bytes_per_pack / pack_factor;
   const int out_vec_size_g = out_vec_size / group_size;
+  // 32 * (tid.y * 2 + simd_gid)
   int out_col = pack_factor * tn * (tid.y * num_simdgroups + simd_gid);
   ws += out_col * bytes_per_pack / pack_factor + simd_lid * out_vec_size_w;
   scales += out_col / group_size + simd_lid * out_vec_size_g;
@@ -552,9 +569,9 @@ METAL_FUNC void fp_qvm_impl(
   if (remaining == 0) {
     for (int i = 0; i < in_vec_size; i += block_size) {
       x_local = *x;
-      scale = dequantize_scale<U>(*scales);
+      scale = dequantize_scale<U, group_size>(*scales);
       w_local = *((device vec_w*)ws);
-      qouter<U, tn * pack_factor>(
+      qouter<U, tn * pack_factor, bits>(
           (thread uint8_t*)&w_local, x_local, scale, result);
 
       x += block_size;
@@ -564,10 +581,10 @@ METAL_FUNC void fp_qvm_impl(
   } else {
     for (int i = block_size; i < in_vec_size; i += block_size) {
       x_local = *x;
-      scale = dequantize_scale<U>(*scales);
+      scale = dequantize_scale<U, group_size>(*scales);
       w_local = *((device vec_w*)ws);
 
-      qouter<U, tn * pack_factor>(
+      qouter<U, tn * pack_factor, bits>(
           (thread uint8_t*)&w_local, x_local, scale, result);
 
       x += block_size;
@@ -576,13 +593,13 @@ METAL_FUNC void fp_qvm_impl(
     }
     if (static_cast<int>(simd_lid) < remaining) {
       x_local = *x;
-      scale = dequantize_scale<U>(*scales);
+      scale = dequantize_scale<U, group_size>(*scales);
       w_local = *((device vec_w*)ws);
     } else {
       x_local = 0;
       scale = 0;
     }
-    qouter<U, tn * pack_factor>(
+    qouter<U, tn * pack_factor, bits>(
         (thread uint8_t*)&w_local, x_local, scale, result);
   }
 
@@ -622,8 +639,7 @@ METAL_FUNC void fp_qmm_t_impl(
     uint3 tid [[threadgroup_position_in_grid]],
     uint lid [[thread_index_in_threadgroup]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
-    uint simd_lid [[thread_index_in_simdgroup]],
-    threadgroup T* lut) {
+    uint simd_lid [[thread_index_in_simdgroup]]) {
   static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE");
   static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE");
 
@@ -631,7 +647,7 @@ METAL_FUNC void fp_qmm_t_impl(
 
   constexpr int WM = 2;
   constexpr int WN = 2;
-  constexpr int pack_factor = get_pack_factor<8>();
+  constexpr int pack_factor = get_pack_factor<8, bits>();
   constexpr int bytes_per_pack = get_bytes_per_pack();
 
   constexpr int BK_padded = (BK + 16 / sizeof(T));
@@ -648,7 +664,8 @@ METAL_FUNC void fp_qmm_t_impl(
       BK_padded,
       1,
       WM * WN * SIMD_SIZE,
-      group_size>;
+      group_size,
+      bits>;
 
   // Set the block
   const int K_w = K * bytes_per_pack / pack_factor;
@@ -667,7 +684,7 @@ METAL_FUNC void fp_qmm_t_impl(
   const short num_els = min(BM, M - y_row);
   const short num_outs = min(BN, N - y_col);
   loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid);
-  loader_w_t loader_w(wl, scales, K, Ws, lut, simd_gid, simd_lid);
+  loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid);
   mma_t mma_op(simd_gid, simd_lid);
 
   if (num_els < BM) {
@@ -746,8 +763,7 @@ METAL_FUNC void fp_qmm_n_impl(
     uint3 tid [[threadgroup_position_in_grid]],
     uint lid [[thread_index_in_threadgroup]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
-    uint simd_lid [[thread_index_in_simdgroup]],
-    threadgroup T* lut) {
+    uint simd_lid [[thread_index_in_simdgroup]]) {
   static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE");
   static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE");
 
@@ -755,7 +771,7 @@ METAL_FUNC void fp_qmm_n_impl(
 
   constexpr int WM = 2;
   constexpr int WN = 2;
-  constexpr int pack_factor = get_pack_factor<8>();
+  constexpr int pack_factor = get_pack_factor<8, bits>();
   constexpr int bytes_per_pack = get_bytes_per_pack();
 
   constexpr int BK_padded = (BK + 16 / sizeof(T));
@@ -773,7 +789,8 @@ METAL_FUNC void fp_qmm_n_impl(
       BN_padded,
       0,
       WM * WN * SIMD_SIZE,
-      group_size>;
+      group_size,
+      bits>;
 
   auto wl = (const device uint8_t*)w;
 
@@ -788,7 +805,7 @@ METAL_FUNC void fp_qmm_n_impl(
   // Make the x loader and mma operation
   const short num_els = min(BM, M - y_row);
   loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid);
-  loader_w_t loader_w(wl, scales, N, Ws, lut, simd_gid, simd_lid);
+  loader_w_t loader_w(wl, scales, N, Ws, simd_gid, simd_lid);
   mma_t mma_op(simd_gid, simd_lid);
 
   if (num_els < BM) {
@@ -1178,7 +1195,6 @@ template <
 
   threadgroup T Xs[BM * BK_padded];
   threadgroup T Ws[BN * BK_padded];
-  threadgroup T lut[16];
 
   if (batched) {
     adjust_matrix_offsets(
@@ -1197,7 +1213,7 @@ template <
         tid);
   }
   fp_qmm_t_impl<T, group_size, bits, aligned_N, BM, BK, BN>(
-      w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid, lut);
+      w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
 }
 
 template <
@@ -1234,7 +1250,6 @@ template <
 
   threadgroup T Xs[BM * BK_padded];
   threadgroup T Ws[BK * BN_padded];
-  threadgroup T lut[16];
 
   if (batched) {
     adjust_matrix_offsets(
@@ -1254,7 +1269,7 @@ template <
   }
 
   fp_qmm_n_impl<T, group_size, bits, BM, BK, BN>(
-      w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid, lut);
+      w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
 }
 
 template <typename T, int group_size, int bits>
@@ -1443,7 +1458,6 @@ template <
 
   threadgroup T Xs[BM * BK_padded];
   threadgroup T Ws[BN * BK_padded];
-  threadgroup T lut[16];
 
   adjust_matrix_offsets(
       x,
@@ -1466,7 +1480,7 @@ template <
       s_strides,
       tid);
   fp_qmm_t_impl<T, group_size, bits, aligned_N, BM, BK, BN>(
-      w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid, lut);
+      w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
 }
 
 template <
@@ -1508,7 +1522,6 @@ template <
 
   threadgroup T Xs[BM * BK_padded];
   threadgroup T Ws[BK * BN_padded];
-  threadgroup T lut[16];
 
   adjust_matrix_offsets(
       x,
@@ -1531,7 +1544,7 @@ template <
       s_strides,
       tid);
   fp_qmm_n_impl<T, group_size, bits, BM, BK, BN>(
-      w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid, lut);
+      w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
 }
 
 template <
@@ -1556,11 +1569,10 @@ template <
     uint3 tid [[threadgroup_position_in_grid]],
     uint simd_group_id [[simdgroup_index_in_threadgroup]],
     uint simd_lane_id [[thread_index_in_simdgroup]]) {
-  constexpr int pack_factor = get_pack_factor<8>();
+  constexpr int pack_factor = get_pack_factor<8, bits>();
   constexpr int bytes_per_pack = get_bytes_per_pack();
   constexpr int BK_padded = (BK + 16 / sizeof(T));
   constexpr int BN_padded = (BN + 16 / sizeof(T));
-  threadgroup T lut[16];
 
   using mma_t = mlx::steel::BlockMMA<
       T,
@@ -1583,7 +1595,8 @@ template <
       transpose ? BK_padded : BN_padded,
       transpose,
       WM * WN * SIMD_SIZE,
-      group_size>;
+      group_size,
+      bits>;
 
   threadgroup T Xs[BM * BK_padded];
   threadgroup T Ws[transpose ? BN * BK_padded : BK * BN_padded];
@@ -1648,7 +1661,6 @@ template <
         scales + index * stride_s,
         transpose ? K : N,
         Ws,
-        lut,
         simd_group_id,
         simd_lane_id);
 

mlx/backend/metal/kernels/fp_quantized.metal

@@ -6,68 +6,68 @@
 #include "mlx/backend/metal/kernels/quantized_utils.h"
 #include "mlx/backend/metal/kernels/fp_quantized.h"
 
-#define instantiate_quantized(mode, name, type) \
+#define instantiate_quantized(mode, name, type, group_size, bits) \
   instantiate_kernel( \
-      #mode "_" #name "_" #type "_gs_32_b_4", \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits, \
       fp_ ## name, \
       type, \
-      32,   \
-      4)
+      group_size,   \
+      bits)
 
-#define instantiate_quantized_batched(mode, name, type, batched) \
+#define instantiate_quantized_batched(mode, name, type, batched, group_size, bits) \
   instantiate_kernel( \
-      #mode  "_" #name "_" #type "_gs_32_b_4_batch_" #batched, \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_batch_" #batched, \
       fp_ ## name,    \
       type,    \
-      32,      \
-      4,       \
+      group_size,      \
+      bits,       \
       batched)
 
-#define instantiate_quantized_aligned(mode, name, type, aligned) \
+#define instantiate_quantized_aligned(mode, name, type, aligned, group_size, bits) \
   instantiate_kernel( \
-      #mode "_" #name "_" #type "_gs_32_b_4_alN_" #aligned, \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_alN_" #aligned, \
       fp_ ## name,    \
       type,    \
-      32,      \
-      4,       \
+      group_size,      \
+      bits,       \
       aligned)
 
-#define instantiate_quantized_aligned_batched(mode, name, type, aligned, batched) \
+#define instantiate_quantized_aligned_batched(mode, name, type, aligned, batched, group_size, bits) \
   instantiate_kernel( \
-      #mode "_" #name "_" #type "_gs_32_b_4_alN_" #aligned "_batch_" #batched, \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_alN_" #aligned "_batch_" #batched, \
       fp_ ## name,    \
       type,    \
-      32,      \
-      4,       \
+      group_size,      \
+      bits,       \
       aligned, \
       batched)
 
-#define instantiate_quantized_quad(mode, name, type, D, batched) \
+#define instantiate_quantized_quad(mode, name, type, D, batched, group_size, bits) \
   instantiate_kernel( \
-      #mode "_" #name "_" #type "_gs_32_b_4_d_" #D "_batch_" #batched, \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_d_" #D "_batch_" #batched, \
       fp_ ## name,    \
       type,    \
-      32,      \
-      4,       \
+      group_size,      \
+      bits,       \
       D,       \
       batched)
 
-#define instantiate_quantized_split_k(mode, name, type, split_k) \
+#define instantiate_quantized_split_k(mode, name, type, split_k, group_size, bits) \
   instantiate_kernel( \
-      #mode "_" #name "_" #type "_gs_32_b_4_spk_" #split_k, \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_spk_" #split_k, \
       fp_ ## name,    \
       type,    \
-      32,      \
-      4,       \
+      group_size,      \
+      bits,       \
       split_k)
 
-#define instantiate_gather_qmm_rhs(func, name, type, bm, bn, bk, wm, wn, transpose) \
+#define instantiate_gather_qmm_rhs(func, name, type, bm, bn, bk, wm, wn, transpose, mode, group_size, bits) \
   instantiate_kernel( \
-      #name "_" #type "_gs_32_b_4_bm_" #bm "_bn_" #bn "_bk_" #bk "_wm_" #wm "_wn_" #wn, \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm_" #bm "_bn_" #bn "_bk_" #bk "_wm_" #wm "_wn_" #wn, \
       func,    \
       type,    \
-      32,      \
-      4,       \
+      group_size,      \
+      bits,       \
       bm,      \
       bn,      \
       bk,      \
@@ -75,43 +75,43 @@
       wn,      \
       transpose)
 
-#define instantiate_quantized_batched_wrap(mode, name, type) \
-  instantiate_quantized_batched(mode, name, type, 1)         \
-  instantiate_quantized_batched(mode, name, type, 0)
+#define instantiate_quantized_batched_wrap(name, type, mode, group_size, bits) \
+  instantiate_quantized_batched(mode, name, type, 1, group_size, bits)         \
+  instantiate_quantized_batched(mode, name, type, 0, group_size, bits)
 
-#define instantiate_quantized_all_batched(type) \
-  instantiate_quantized_batched_wrap(mxfp4, qmv_fast, type) \
-  instantiate_quantized_batched_wrap(mxfp4, qmv, type)      \
-  instantiate_quantized_batched_wrap(mxfp4, qvm, type)      \
-  instantiate_quantized_batched_wrap(mxfp4, qmm_n, type)
+#define instantiate_quantized_all_batched(type, mode, group_size, bits) \
+  instantiate_quantized_batched_wrap(qmv_fast, type, mode, group_size, bits) \
+  instantiate_quantized_batched_wrap(qmv, type, mode, group_size, bits)      \
+  instantiate_quantized_batched_wrap(qvm, type, mode, group_size, bits) \
+  instantiate_quantized_batched_wrap(qmm_n, type, mode, group_size, bits)
 
-#define instantiate_quantized_all_single(type) \
-  instantiate_quantized(mxfp4, gather_qmv_fast, type) \
-  instantiate_quantized(mxfp4, gather_qmv, type)      \
-  instantiate_quantized(mxfp4, gather_qvm, type)      \
-  instantiate_quantized(mxfp4, gather_qmm_n, type)
+#define instantiate_quantized_all_single(type, mode, group_size, bits) \
+  instantiate_quantized(mode, gather_qmv_fast, type, group_size, bits) \
+  instantiate_quantized(mode, gather_qmv, type, group_size, bits)      \
+  instantiate_quantized(mode, gather_qvm, type, group_size, bits) \
+  instantiate_quantized(mode, gather_qmm_n, type, group_size, bits)
 
-#define instantiate_quantized_all_aligned(type) \
-  instantiate_quantized_aligned(mxfp4, gather_qmm_t, type, true)      \
-  instantiate_quantized_aligned(mxfp4, gather_qmm_t, type, false)     \
-  instantiate_quantized_aligned_batched(mxfp4, qmm_t, type, true, 1)  \
-  instantiate_quantized_aligned_batched(mxfp4, qmm_t, type, true, 0)  \
-  instantiate_quantized_aligned_batched(mxfp4, qmm_t, type, false, 1) \
-  instantiate_quantized_aligned_batched(mxfp4, qmm_t, type, false, 0)
+#define instantiate_quantized_all_aligned(type, mode, group_size, bits) \
+  instantiate_quantized_aligned(mode, gather_qmm_t, type, true, group_size, bits)      \
+  instantiate_quantized_aligned(mode, gather_qmm_t, type, false, group_size, bits)     \
+  instantiate_quantized_aligned_batched(mode, qmm_t, type, true, 1, group_size, bits)  \
+  instantiate_quantized_aligned_batched(mode, qmm_t, type, true, 0, group_size, bits)  \
+  instantiate_quantized_aligned_batched(mode, qmm_t, type, false, 1, group_size, bits) \
+  instantiate_quantized_aligned_batched(mode, qmm_t, type, false, 0, group_size, bits)
 
-#define instantiate_quantized_all_quad(type) \
-  instantiate_quantized_quad(mxfp4, qmv_quad, type, 64, 1)  \
-  instantiate_quantized_quad(mxfp4, qmv_quad, type, 64, 0)  \
-  instantiate_quantized_quad(mxfp4, qmv_quad, type, 128, 1) \
-  instantiate_quantized_quad(mxfp4, qmv_quad, type, 128, 0)
+#define instantiate_quantized_all_quad(type, mode, group_size, bits) \
+  instantiate_quantized_quad(mode, qmv_quad, type, 64, 1, group_size, bits)  \
+  instantiate_quantized_quad(mode, qmv_quad, type, 64, 0, group_size, bits)  \
+  instantiate_quantized_quad(mode, qmv_quad, type, 128, 1, group_size, bits) \
+  instantiate_quantized_quad(mode, qmv_quad, type, 128, 0, group_size, bits)
 
-#define instantiate_quantized_all_splitk(type) \
-  instantiate_quantized_split_k(mxfp4, qvm_split_k, type, 8) \
-  instantiate_quantized_split_k(mxfp4, qvm_split_k, type, 32)
+#define instantiate_quantized_all_splitk(type, mode, group_size, bits) \
+  instantiate_quantized_split_k(mode, qvm_split_k, type, 8, group_size, bits) \
+  instantiate_quantized_split_k(mode, qvm_split_k, type, 32, group_size, bits)
 
-#define instantiate_quantized_all_rhs(type) \
-  instantiate_gather_qmm_rhs(fp_gather_qmm_rhs, mxfp4_gather_qmm_rhs_nt, type, 16, 32, 32, 1, 2, true) \
-  instantiate_gather_qmm_rhs(fp_gather_qmm_rhs, mxfp4_gather_qmm_rhs_nn, type, 16, 32, 32, 1, 2, false)
+#define instantiate_quantized_all_rhs(type, mode, group_size, bits) \
+  instantiate_gather_qmm_rhs(fp_gather_qmm_rhs, gather_qmm_rhs_nt, type, 16, 32, 32, 1, 2, true, mode, group_size, bits) \
+  instantiate_gather_qmm_rhs(fp_gather_qmm_rhs, gather_qmm_rhs_nn, type, 16, 32, 32, 1, 2, false, mode, group_size, bits)
 
 #define instantiate_quantize_dequantize(type, mode, group_size, bits) \
   instantiate_kernel( \
@@ -127,19 +127,19 @@
     group_size,  \
     bits)
 
-#define instantiate_quantize_dequantize_modes(type) \
-  instantiate_quantize_dequantize(type, mxfp4, 32, 4) \
-  instantiate_quantize_dequantize(type, nvfp4, 16, 4) \
-  instantiate_quantize_dequantize(type, mxfp8, 32, 8)
+#define instantiate_quantized_modes(type, mode, group_size, bits) \
+  instantiate_quantized_all_batched(type, mode, group_size, bits) \
+  instantiate_quantized_all_single(type, mode, group_size, bits)  \
+  instantiate_quantized_all_quad(type, mode, group_size, bits)    \
+  instantiate_quantized_all_splitk(type, mode, group_size, bits)  \
+  instantiate_quantized_all_aligned(type, mode, group_size, bits) \
+  instantiate_quantized_all_rhs(type, mode, group_size, bits)     \
+  instantiate_quantize_dequantize(type, mode, group_size, bits)
 
 #define instantiate_quantized_types(type) \
-  instantiate_quantized_all_batched(type) \
-  instantiate_quantized_all_quad(type)    \
-  instantiate_quantized_all_splitk(type)  \
-  instantiate_quantized_all_single(type)  \
-  instantiate_quantized_all_aligned(type) \
-  instantiate_quantized_all_rhs(type)     \
-  instantiate_quantize_dequantize_modes(type)
+  instantiate_quantized_modes(type, nvfp4, 16, 4) \
+  instantiate_quantized_modes(type, mxfp8, 32, 8) \
+  instantiate_quantized_modes(type, mxfp4, 32, 4)
 
 instantiate_quantized_types(float)
 instantiate_quantized_types(bfloat16_t)

mlx/backend/metal/kernels/fp_quantized_nax.h

@@ -17,9 +17,9 @@ using namespace metal;
 MLX_MTL_CONST int SIMD_SIZE = 32;
 MLX_MTL_CONST int QUAD_SIZE = 4;
 
-template <int wsize = 8>
+template <int wsize = 8, int bits>
 inline constexpr short get_pack_factor() {
-  return wsize / 4;
+  return wsize / bits;
 }
 
 template <int wsize = 8>
@@ -27,15 +27,20 @@ inline constexpr short get_bytes_per_pack() {
   return wsize / 8;
 }
 
-template <typename T>
+template <typename T, int group_size>
 static inline T dequantize_scale(uint8_t s) {
-  return T(*(thread fp8_e8m0*)(&s));
+  if constexpr (group_size == 16) {
+    // Use nv scale
+    return T(*(thread fp8_e4m3*)(&s));
+  } else {
+    return T(*(thread fp8_e8m0*)(&s));
+  }
 }
 
 template <int bits>
 struct Quantize {
   uint8_t operator()(float x) {
-    if constexpr (bits == 8) {
+    if (bits == 8) {
       return fp8_e4m3(x).bits;
     } else {
       return fp4_e2m1(x).bits;
@@ -43,26 +48,24 @@ struct Quantize {
   }
 };
 
-template <int bits>
+template <int bits, typename U = float>
 struct Dequantize {
-  float operator()(uint8_t x) {
+  U operator()(uint8_t x) {
     if constexpr (bits == 8) {
-      return float(*(thread fp8_e4m3*)(&x));
+      return U(*(thread fp8_e4m3*)(&x));
     } else {
-      return float(*(thread fp4_e2m1*)(&x));
+      return U(*(thread fp4_e2m1*)(&x));
     }
   }
 };
 
-template <typename U, int N>
-inline void dequantize(
-    const device uint8_t* w,
-    U scale,
-    threadgroup U* w_local,
-    const threadgroup U* lut) {
-  for (int i = 0; i < (N / 2); i++) {
-    w_local[2 * i] = scale * lut[w[i] & 0xf];
-    w_local[2 * i + 1] = scale * lut[(w[i] >> 4) & 0xf];
+template <typename U, int bits>
+inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) {
+  if constexpr (bits == 4) {
+    w_local[0] = scale * Dequantize<4, U>{}(w);
+    w_local[1] = scale * Dequantize<4, U>{}(w >> 4);
+  } else {
+    w_local[0] = scale * Dequantize<8, U>{}(w);
   }
 }
 
@@ -73,22 +76,21 @@ template <
     short dst_ld,
     short reduction_dim,
     short tgp_size,
-    short group_size>
+    short group_size,
+    short bits>
 struct QuantizedBlockLoader {
-  static_assert(
-      BCOLS % group_size == 0,
-      "The group size should be divisible by the columns");
-
-  MLX_MTL_CONST short pack_factor = get_pack_factor<8>();
+  MLX_MTL_CONST short pack_factor = get_pack_factor<8, bits>();
   MLX_MTL_CONST short bytes_per_pack = get_bytes_per_pack();
   MLX_MTL_CONST short BCOLS_PACKED = BCOLS / pack_factor;
   MLX_MTL_CONST short n_reads =
       (BCOLS_PACKED * BROWS < tgp_size) ? 1 : (BCOLS_PACKED * BROWS) / tgp_size;
-  MLX_MTL_CONST short n_groups = BCOLS / group_size;
 
-  static_assert(
-      (BCOLS_PACKED / n_reads) == n_groups,
-      "Other configurations are not yet supported");
+  MLX_MTL_CONST short n_reads_per_scale = (n_reads * pack_factor) <= group_size
+      ? n_reads
+      : (group_size / pack_factor);
+  MLX_MTL_CONST short n_steps_per_read = n_reads / n_reads_per_scale;
+
+  MLX_MTL_CONST short n_groups = BCOLS / group_size;
 
   const int src_ld;
   const int tile_stride;
@@ -103,14 +105,12 @@ struct QuantizedBlockLoader {
   threadgroup T* dst;
   const device uint8_t* src;
   const device uint8_t* scales;
-  threadgroup T* lut;
 
   QuantizedBlockLoader(
       const device uint8_t* src_,
       const device uint8_t* scales_,
       const int src_ld_,
       threadgroup T* dst_,
-      threadgroup T* lut_,
       ushort simd_group_id [[simdgroup_index_in_threadgroup]],
       ushort simd_lane_id [[thread_index_in_simdgroup]])
       : src_ld(src_ld_),
@@ -125,23 +125,21 @@ struct QuantizedBlockLoader {
         dst(dst_ + bi * dst_ld + bj * pack_factor),
         src(src_ + bi * src_ld * bytes_per_pack / pack_factor +
             bj * bytes_per_pack),
-        scales(scales_ + bi * src_ld / group_size + group_id),
-        lut(lut_) {
-    if (simd_group_id == 0 && simd_lane_id < 16) {
-      lut[simd_lane_id] = static_cast<T>(FP4_LUT[simd_lane_id]);
-    }
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-  }
+        scales(scales_ + bi * src_ld / group_size + group_id) {}
 
   void load_unsafe() const {
     if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) {
       return;
     }
 
-    T scale = dequantize_scale<T>(*scales);
-    for (int i = 0; i < n_reads; i++) {
-      dequantize<T, pack_factor>(
-          src + i * bytes_per_pack, scale, dst + i * pack_factor, lut);
+    int k = 0;
+    for (int i = 0; i < n_steps_per_read; i++) {
+      T scale = dequantize_scale<T, group_size>(scales[i]);
+      for (int j = 0; j < n_reads_per_scale; j++) {
+        dequantize<T, bits>(
+            src[k * bytes_per_pack], scale, dst + k * pack_factor);
+        k++;
+      }
     }
   }
 
@@ -164,28 +162,21 @@ struct QuantizedBlockLoader {
       return;
     }
 
-    T scale = dequantize_scale<T>(*scales);
-    for (int i = 0; i < n_reads; i++) {
-      dequantize<T, pack_factor>(
-          (device uint8_t*)(src + i * bytes_per_pack),
-          scale,
-          dst + i * pack_factor,
-          lut);
+    int k = 0;
+    for (int i = 0; i < n_steps_per_read; i++) {
+      T scale = dequantize_scale<T, group_size>(scales[i]);
+      for (int j = 0; j < n_reads_per_scale; j++) {
+        dequantize<T, bits>(
+            src[k * bytes_per_pack], scale, dst + k * pack_factor);
+        k++;
+      }
     }
   }
 
   void next() {
     src += tile_stride;
     if (reduction_dim == 1) {
-      // if (group_steps > 1) {
-      //   group_step_cnt++;
-      //   if (group_step_cnt == group_steps) {
-      //     group_step_cnt = 0;
-      //     scales++;
-      //   }
-      // } else {
       scales += n_groups;
-      // }
     } else {
       scales += n_groups * group_stride;
     }
@@ -217,14 +208,13 @@ METAL_FUNC void fp_qmm_t_impl(
     uint3 tid [[threadgroup_position_in_grid]],
     uint lid [[thread_index_in_threadgroup]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
-    uint simd_lid [[thread_index_in_simdgroup]],
-    threadgroup Wtype* lut) {
+    uint simd_lid [[thread_index_in_simdgroup]]) {
   static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE");
   static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE");
 
   (void)lid;
 
-  constexpr int pack_factor = get_pack_factor<8>();
+  constexpr int pack_factor = get_pack_factor<8, bits>();
   constexpr int bytes_per_pack = get_bytes_per_pack();
 
   constexpr int BK_padded = (BK + 16 / sizeof(Wtype));
@@ -237,7 +227,8 @@ METAL_FUNC void fp_qmm_t_impl(
       BK_padded,
       1,
       WM * WN * SIMD_SIZE,
-      group_size>;
+      group_size,
+      bits>;
 
   // Set the block
   const int K_w = K * bytes_per_pack / pack_factor;
@@ -253,7 +244,7 @@ METAL_FUNC void fp_qmm_t_impl(
   y += y_row * static_cast<int64_t>(N) + y_col;
 
   // Make the weight loader
-  loader_w_t loader_w(wl, scales, K, Ws, lut, simd_gid, simd_lid);
+  loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid);
 
   constexpr short UM = 16;
   constexpr short UN = 32;
@@ -362,7 +353,6 @@ METAL_FUNC void fp_qmm_n_impl(
     const device uint8_t* scales,
     const device T* x,
     device T* y,
-    threadgroup T* Xs,
     threadgroup T* Ws,
     const constant int& K,
     const constant int& N,
@@ -370,24 +360,18 @@ METAL_FUNC void fp_qmm_n_impl(
     uint3 tid [[threadgroup_position_in_grid]],
     uint lid [[thread_index_in_threadgroup]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
-    uint simd_lid [[thread_index_in_simdgroup]],
-    threadgroup T* lut) {
+    uint simd_lid [[thread_index_in_simdgroup]]) {
   static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE");
   static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE");
 
   (void)lid;
+  (void)M;
 
-  constexpr int pack_factor = get_pack_factor<8>();
+  constexpr int pack_factor = get_pack_factor<8, bits>();
   constexpr int bytes_per_pack = get_bytes_per_pack();
 
-  constexpr int BK_padded = (BK + 16 / sizeof(T));
   constexpr int BN_padded = (BN + 16 / sizeof(T));
 
-  // Instantiate the appropriate BlockMMA and Loader
-  using mma_t = mlx::steel::
-      BlockMMA<T, T, BM, BN, BK, WM, WN, false, false, BK_padded, BN_padded>;
-  using loader_x_t = mlx::steel::
-      BlockLoader<T, BM, BK, BK_padded, 1, WM * WN * SIMD_SIZE, 1, 4>;
   using loader_w_t = QuantizedBlockLoader<
       T,
       BK,
@@ -395,91 +379,91 @@ METAL_FUNC void fp_qmm_n_impl(
       BN_padded,
       0,
       WM * WN * SIMD_SIZE,
-      group_size>;
+      group_size,
+      bits>;
+
+  // Set the block
+  const int K_w = K * bytes_per_pack / pack_factor;
+  const int K_g = K / group_size;
+  const int y_row = tid.y * BM;
+  const int y_col = tid.x * BN;
 
   auto wl = (const device uint8_t*)w;
 
-  // Set the block
-  const int y_row = tid.y * BM;
-  const int y_col = tid.x * BN;
   x += y_row * static_cast<int64_t>(K);
-  wl += y_col * bytes_per_pack / pack_factor;
-  scales += y_col / group_size;
+  wl += y_col * K_w;
+  scales += y_col * K_g;
   y += y_row * static_cast<int64_t>(N) + y_col;
 
   // Make the x loader and mma operation
-  const short num_els = min(BM, M - y_row);
-  loader_x_t loader_x(x, K, Xs, simd_gid, simd_lid);
-  loader_w_t loader_w(wl, scales, N, Ws, lut, simd_gid, simd_lid);
-  mma_t mma_op(simd_gid, simd_lid);
+  // const short num_els = min(BM, M - y_row);
+  // const short num_outs = min(BN, N - y_col);
+  loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid);
 
-  if (num_els < BM) {
-    if ((K % BK) != 0) {
-      const int k_blocks = K / BK;
-      for (int k = 0; k < k_blocks; k++) {
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        loader_x.load_safe(short2(BK, num_els));
-        loader_w.load_unsafe();
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        mma_op.mma(Xs, Ws);
-        loader_x.next();
-        loader_w.next();
-      }
-      const short num_k = K - k_blocks * BK;
-      threadgroup_barrier(mem_flags::mem_threadgroup);
-      loader_x.load_safe(short2(num_k, num_els));
-      loader_w.load_safe(short2(BN, num_k));
-      threadgroup_barrier(mem_flags::mem_threadgroup);
-      mma_op.mma(Xs, Ws);
-    } else {
-      for (int k = 0; k < K; k += BK) {
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        loader_x.load_safe(short2(BK, num_els));
-        loader_w.load_unsafe();
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        mma_op.mma(Xs, Ws);
-        loader_x.next();
-        loader_w.next();
-      }
-    }
-  } else {
-    if ((K % BK) != 0) {
-      const int k_blocks = K / BK;
-      for (int k = 0; k < k_blocks; k++) {
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        loader_x.load_unsafe();
-        loader_w.load_unsafe();
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        mma_op.mma(Xs, Ws);
-        loader_x.next();
-        loader_w.next();
-      }
-      const short num_k = K - k_blocks * BK;
-      threadgroup_barrier(mem_flags::mem_threadgroup);
-      loader_x.load_safe(short2(num_k, BM));
-      loader_w.load_safe(short2(BN, num_k));
-      threadgroup_barrier(mem_flags::mem_threadgroup);
-      mma_op.mma(Xs, Ws);
-    } else {
-      for (int k = 0; k < K; k += BK) {
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        loader_x.load_unsafe();
-        loader_w.load_unsafe();
-        threadgroup_barrier(mem_flags::mem_threadgroup);
-        mma_op.mma(Xs, Ws);
-        loader_x.next();
-        loader_w.next();
-      }
+  constexpr short UM = 16;
+  constexpr short UN = 32;
+  constexpr short UK = 16;
+  constexpr short SM = BM / WM;
+  constexpr short SN = BN / WN;
+  constexpr short SK = 32;
+
+  constexpr short TM = SM / UM;
+  constexpr short TN = SN / UN;
+  constexpr short TK = SK / UK;
+
+  const short tm = SM * (simd_gid / WN);
+  const short tn = SN * (simd_gid % WN);
+
+  const short ldb_tgp = BN_padded;
+
+  constexpr bool transpose_a = false;
+  constexpr bool transpose_b = false;
+
+  using AccumType = float;
+
+  using ASubTile = NAXSubTile<T, UM, UK>;
+  using BSubTile = NAXSubTile<T, UK, UN>;
+  using DSubTile = NAXSubTile<AccumType, UM, UN>;
+
+  NAXTile<AccumType, TM, TN, DSubTile> Dtile;
+
+  Dtile.clear();
+
+  x += tm * K;
+
+  for (int k = 0; k < K; k += BK) {
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    loader_w.load_unsafe();
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    STEEL_PRAGMA_NO_UNROLL
+    for (int kk1 = 0; kk1 < BK; kk1 += SK) {
+      NAXTile<T, TM, TK, ASubTile> Atile;
+      NAXTile<Wtype, TK, TN, BSubTile> Btile;
+
+      volatile int compiler_barrier;
+
+      Atile.load(x + kk1, K);
+      Btile.template load<T, BN_padded, 1>(Ws + tn + kk1 * ldb_tgp);
+
+      tile_matmad_nax(
+          Dtile,
+          Atile,
+          metal::bool_constant<transpose_a>{},
+          Btile,
+          metal::bool_constant<transpose_b>{});
+
+      (void)compiler_barrier;
     }
+
+    x += BK;
+    loader_w.next();
   }
 
   // Store results to device memory
   threadgroup_barrier(mem_flags::mem_threadgroup);
-  if (num_els < BM) {
-    mma_op.store_result_safe(y, N, short2(BN, num_els));
-  } else {
-    mma_op.store_result(y, N);
-  }
+
+  Dtile.store(y + tm * N + tn, N);
 }
 
 template <typename T, typename S>
@@ -603,7 +587,6 @@ template <
   constexpr int BK_padded = (BK + 16 / sizeof(Wtype));
 
   threadgroup Wtype Ws[BN * BK_padded];
-  threadgroup Wtype lut[16];
 
   if (batched) {
     adjust_matrix_offsets(
@@ -622,7 +605,7 @@ template <
         tid);
   }
   fp_qmm_t_impl<T, group_size, bits, aligned_N, BM, BK, BN, WM, WN, Wtype>(
-      w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid, lut);
+      w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
 }
 
 template <
@@ -662,7 +645,6 @@ template <
 
   threadgroup T Xs[BM * BK_padded];
   threadgroup T Ws[BK * BN_padded];
-  threadgroup T lut[16];
 
   if (batched) {
     adjust_matrix_offsets(
@@ -682,7 +664,7 @@ template <
   }
 
   fp_qmm_n_impl<T, group_size, bits, BM, BK, BN, WM, WN, Wtype>(
-      w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid, lut);
+      w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
 }
 
 template <
@@ -726,7 +708,6 @@ template <
   constexpr int BK_padded = (BK + 16 / sizeof(Wtype));
 
   threadgroup Wtype Ws[BN * BK_padded];
-  threadgroup Wtype lut[16];
 
   adjust_matrix_offsets(
       x,
@@ -749,7 +730,7 @@ template <
       s_strides,
       tid);
   fp_qmm_t_impl<T, group_size, bits, aligned_N, BM, BK, BN, WM, WN, Wtype>(
-      w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid, lut);
+      w, scales, x, y, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
 }
 
 template <
@@ -794,7 +775,6 @@ template <
 
   threadgroup T Xs[BM * BK_padded];
   threadgroup T Ws[BK * BN_padded];
-  threadgroup T lut[16];
 
   adjust_matrix_offsets(
       x,
@@ -817,7 +797,7 @@ template <
       s_strides,
       tid);
   fp_qmm_n_impl<T, group_size, bits, BM, BK, BN, WM, WN, Wtype>(
-      w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid, lut);
+      w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid);
 }
 
 template <
@@ -843,13 +823,11 @@ template <
     uint3 tid [[threadgroup_position_in_grid]],
     uint simd_group_id [[simdgroup_index_in_threadgroup]],
     uint simd_lane_id [[thread_index_in_simdgroup]]) {
-  constexpr int pack_factor = get_pack_factor<8>();
+  constexpr int pack_factor = get_pack_factor<8, bits>();
   constexpr int bytes_per_pack = get_bytes_per_pack();
   constexpr int BK_padded = (BK + 16 / sizeof(Wtype));
   constexpr int BN_padded = (BN + 16 / sizeof(Wtype));
 
-  threadgroup Wtype lut[16];
-
   using loader_w_t = QuantizedBlockLoader<
       Wtype,
       transpose ? BN : BK,
@@ -857,7 +835,8 @@ template <
       transpose ? BK_padded : BN_padded,
       transpose,
       WM * WN * SIMD_SIZE,
-      group_size>;
+      group_size,
+      bits>;
 
   threadgroup Wtype Ws[transpose ? BN * BK_padded : BK * BN_padded];
 
@@ -954,7 +933,6 @@ template <
         scales + index * stride_s,
         transpose ? K : N,
         Ws,
-        lut,
         simd_group_id,
         simd_lane_id);