Use uv

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

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

@@ -21,4 +21,11 @@ runs:
         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 ${{ inputs.arch }}
+
+        auditwheel repair dist/mlx_cuda*.whl \
+          --plat manylinux_2_35_${{ inputs.arch }} \
+          --exclude libcublas* \
+          --exclude libcuda* \
+          --exclude libcudnn* \
+          --exclude libnccl* \
+          --exclude libnvrtc*

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

@@ -18,19 +18,21 @@ inputs:
 runs:
   using: "composite"
   steps:
-    - name: Generate package stubs
+    - name: Build MLX
       shell: bash
-      run: |
-        pip install -e ".[dev]" -v
-        pip install typing_extensions
-        python setup.py generate_stubs
+      run: pip install -e . -v
+
     - name: Build Python package
       shell: bash
       run: |
         pip install auditwheel patchelf build
         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

@@ -21,19 +21,13 @@ runs:
         if ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }} ; then
           # There is no GPU in arm64 runner, use a common arch.
           CMAKE_ARGS="$CMAKE_ARGS -DMLX_CUDA_ARCHITECTURES=90a"
-          # Can not build tests when the built executables can not run.
-          CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF"
+          # 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: Generate package stubs
-      shell: sh
-      run: |
-        pip install typing_extensions
-        python setup.py generate_stubs
-
     - name: Build CPP only
       shell: bash
       run: |

.github/actions/build-macos/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.10.2
         pip install -e . -v
 
-    - name: Generate package stubs
-      shell: bash -l {0}
-      run: |
-        pip install typing_extensions
-        python setup.py generate_stubs
-
     - name: Install tests dependencies
       shell: bash -l {0}
       run: |

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

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

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

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

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

.github/workflows/nightly.yml

@@ -65,6 +65,11 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
       - uses: ./.github/actions/build-macos
+      - name: Build macOS 26 package
+        uses: ./.github/actions/build-macos-release
+        with:
+          macos-target: 26.0
+          build-backend: ${{ matrix.python-version == '3.10' }}
       - name: Build macOS 15 package
         uses: ./.github/actions/build-macos-release
         with:
@@ -88,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@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: ubuntu-22.04
     steps:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/build-docs
-    
+
   deploy_documentation:
+    if: inputs.publish
     needs: build_documentation
     permissions:
       pages: write
@@ -51,7 +52,7 @@ jobs:
     runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }}
     env:
       PYPI_RELEASE: 1
-      DEV_RELEASE: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}
+      DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }}
     steps:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/setup-linux
@@ -60,7 +61,7 @@ jobs:
           use-ccache: false
       - uses: ./.github/actions/build-linux-release
         with:
-          build-backend: ${{ matrix.python-version == '3.10' }}
+          build-backend: ${{ matrix.python_version == '3.10' }}
           arch: ${{ matrix.arch }}
       - name: Upload MLX artifacts
         uses: actions/upload-artifact@v6
@@ -68,6 +69,7 @@ jobs:
           overwrite: true
           name: linux-wheels-${{ matrix.python_version }}-${{ matrix.arch }}
           path: wheelhouse/mlx-*.whl
+          if-no-files-found: error
       - name: Upload CPU artifacts
         if: matrix.python_version == '3.10'
         uses: actions/upload-artifact@v6
@@ -75,7 +77,8 @@ jobs:
           overwrite: true
           name: mlx-cpu-${{ matrix.arch }}
           path: wheelhouse/mlx_cpu-*.whl
-  
+          if-no-files-found: error
+
   build_mac_release:
     if: github.repository == 'ml-explore/mlx'
     strategy:
@@ -84,8 +87,7 @@ jobs:
     runs-on: [self-hosted, macos]
     env:
       PYPI_RELEASE: 1
-      DEV_RELEASE: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}
-
+      DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }}
     steps:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/setup-macos
@@ -96,13 +98,8 @@ jobs:
         shell: bash -l {0}
         run: |
           pip install --upgrade pip
-          pip install cmake setuptools nanobind==2.10.2
+          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:
@@ -119,6 +116,7 @@ jobs:
           overwrite: true
           name: mac-wheels-${{ matrix.python-version }}
           path: dist/mlx-*.whl
+          if-no-files-found: error
       - name: Upload Metal artifacts
         if: matrix.python-version == '3.10'
         uses: actions/upload-artifact@v6
@@ -126,6 +124,7 @@ jobs:
           overwrite: true
           name: mlx-metal
           path: dist/mlx_metal-*.whl
+          if-no-files-found: error
 
   build_cuda_release:
     if: github.repository == 'ml-explore/mlx'
@@ -136,7 +135,7 @@ jobs:
     runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22-large' || 'ubuntu-22-large-arm' }}
     env:
       PYPI_RELEASE: 1
-      DEV_RELEASE: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}
+      DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }}
     steps:
       - uses: actions/checkout@v6
       - uses: ./.github/actions/setup-linux
@@ -151,17 +150,18 @@ jobs:
         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@v7
@@ -175,29 +175,32 @@ jobs:
           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@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/
@@ -205,11 +208,11 @@ 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@v7
@@ -218,8 +221,9 @@ jobs:
           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/
@@ -227,11 +231,11 @@ 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@v7
@@ -239,8 +243,9 @@ jobs:
           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)
@@ -276,11 +339,13 @@ if(MLX_BUILD_PYTHON_BINDINGS)
     Python 3.10
     COMPONENTS Interpreter Development.Module
     REQUIRED)
-  execute_process(
-    COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
-    OUTPUT_STRIP_TRAILING_WHITESPACE
-    OUTPUT_VARIABLE nanobind_ROOT)
-  find_package(nanobind CONFIG REQUIRED)
+  FetchContent_Declare(
+    nanobind
+    GIT_REPOSITORY https://github.com/wjakob/nanobind.git
+    GIT_TAG v2.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

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

examples/python/linear_regression.py

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

examples/python/logistic_regression.py

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

mlx/array.cpp

@@ -96,7 +96,7 @@ array::array(
     deleter(data);
   } else {
     auto wrapped_deleter = [deleter](allocator::Buffer buffer) {
-      auto ptr = buffer.ptr();
+      auto ptr = buffer.raw_ptr();
       allocator::release(buffer);
       return deleter(ptr);
     };

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

mlx/backend/cuda/CMakeLists.txt

@@ -29,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
@@ -156,6 +157,18 @@ 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.
@@ -214,3 +227,14 @@ target_link_libraries(mlx PRIVATE cudnn_frontend)
 # Link with the actual cuDNN libraries.
 include(${cudnn_frontend_SOURCE_DIR}/cmake/cuDNN.cmake)
 target_link_libraries(mlx PRIVATE CUDNN::cudnn_all)
+
+# Use header-only CUTLASS.
+FetchContent_Declare(
+  cutlass
+  GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
+  GIT_TAG v4.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/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,15 +134,38 @@ 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 = 8;
             auto dim0 = ndim > 0 ? shape.back() : 1;

mlx/backend/cuda/cublas_utils.cpp

@@ -40,21 +40,6 @@ cublasLtMatmulPreference_t get_preference(cu::Device& device) {
   return pref.pref_;
 }
 
-void* allocate_workspace(cu::CommandEncoder& encoder, size_t workspace_size) {
-  if (workspace_size == 0) {
-    return nullptr;
-  }
-
-  // Ensure workspace is 256-byte aligned
-  int nbytes = cuda::ceil_div(workspace_size, 256) * 256;
-  array workspace(
-      cu::malloc_async(nbytes, encoder),
-      {static_cast<int>(workspace_size)},
-      int8);
-  encoder.add_temporary(workspace);
-  return gpu_ptr<void>(workspace);
-}
-
 cublasLtMatrixLayout_t create_matrix_layout(
     cudaDataType_t type,
     uint64_t rows,
@@ -193,8 +178,7 @@ void CublasMatmulBase::execute_matmul(
     }
   }
 
-  void* workspace_ptr =
-      cublas_utils::allocate_workspace(encoder, heuristic_.workspaceSize);
+  void* workspace_ptr = allocate_workspace(encoder, heuristic_.workspaceSize);
 
   // Execute matmul
   auto capture = encoder.capture_context();

mlx/backend/cuda/cublas_utils.h

@@ -12,8 +12,6 @@ namespace cublas_utils {
 // Get the shared cublas preference for a device
 cublasLtMatmulPreference_t get_preference(cu::Device& device);
 
-void* allocate_workspace(cu::CommandEncoder& encoder, size_t workspace_size);
-
 cublasLtMatrixLayout_t create_matrix_layout(
     cudaDataType_t type,
     uint64_t rows,

mlx/backend/cuda/cudnn_utils.cpp

@@ -94,15 +94,7 @@ fe::error_t DnnGraph::encode_capturing(
 void* DnnGraph::prepare_workspace(cu::CommandEncoder& encoder) {
   int64_t workspace_size = 0;
   CHECK_CUDNN_FE_ERROR(get_workspace_size(workspace_size));
-  if (workspace_size > 0) {
-    array workspace(
-        cu::malloc_async(workspace_size, encoder),
-        {static_cast<int>(workspace_size)},
-        uint8);
-    encoder.add_temporary(workspace);
-    return gpu_ptr<void>(workspace);
-  }
-  return nullptr;
+  return allocate_workspace(encoder, workspace_size);
 }
 
 void DnnGraph::set_tensor_attrs(

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

@@ -33,7 +33,6 @@ void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
 
 NO_GPU(BlockMaskedMM)
 NO_GPU(FFT)
-NO_GPU(GatherMM)
 NO_GPU(GatherQMM)
 NO_GPU(Hadamard)
 NO_GPU_MULTI(LUF)

mlx/backend/cuda/utils.cpp

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

@@ -43,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/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)

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/kernels/CMakeLists.txt

@@ -6,6 +6,7 @@ set(BASE_HEADERS
     erf.h
     expm1f.h
     fp8.h
+    logging.h
     utils.h)
 
 function(build_kernel_base TARGET SRCFILE DEPS)
@@ -20,6 +21,9 @@ function(build_kernel_base TARGET SRCFILE DEPS)
   if(MLX_METAL_DEBUG)
     set(METAL_FLAGS ${METAL_FLAGS} -gline-tables-only -frecord-sources)
   endif()
+  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}
                     "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")

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;
@@ -369,15 +360,14 @@ 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 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;
   (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 BN_padded = (BN + 16 / sizeof(T));
@@ -389,7 +379,8 @@ 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;
@@ -407,7 +398,7 @@ METAL_FUNC void fp_qmm_n_impl(
   // Make the x loader and mma operation
   // 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, 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;
@@ -596,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(
@@ -615,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 <
@@ -655,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(
@@ -675,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 <
@@ -719,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,
@@ -742,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 <
@@ -787,7 +775,6 @@ template <
 
   threadgroup T Xs[BM * BK_padded];
   threadgroup T Ws[BK * BN_padded];
-  threadgroup T lut[16];
 
   adjust_matrix_offsets(
       x,
@@ -810,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 <
@@ -836,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,
@@ -850,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];
 
@@ -947,7 +933,6 @@ template <
         scales + index * stride_s,
         transpose ? K : N,
         Ws,
-        lut,
         simd_group_id,
         simd_lane_id);
 

mlx/backend/metal/kernels/fp_quantized_nax.metal

@@ -8,41 +8,41 @@
 #include "mlx/backend/metal/kernels/fp_quantized_nax.h"
 
 
-#define instantiate_quantized_batched(mode, name, type, bm, bn, bk, wm, wn, batched) \
+#define instantiate_quantized_batched(mode, name, type, bm, bn, bk, wm, wn, batched, group_size, bits) \
   instantiate_kernel( \
-      #mode "_" #name "_" #type "_gs_32_b_4_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_batch_" #batched, \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_batch_" #batched, \
       fp_ ## name,  \
       type,         \
-      32,           \
-      4,            \
+      group_size,           \
+      bits,            \
       batched)
 
-#define instantiate_quantized_aligned(mode, name, type, bm, bn, bk, wm, wn, aligned) \
+#define instantiate_quantized_aligned(mode, name, type, bm, bn, bk, wm, wn, aligned, group_size, bits) \
   instantiate_kernel( \
-      #mode "_" #name "_" #type "_gs_32_b_4_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_alN_" #aligned, \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_alN_" #aligned, \
       fp_ ## name, \
       type,        \
-      32,          \
-      4,           \
+      group_size,          \
+      bits,           \
       aligned)
 
-#define instantiate_quantized_aligned_batched(mode, name, type, bm, bn, bk, wm, wn, aligned, batched) \
+#define instantiate_quantized_aligned_batched(mode, name, type, bm, bn, bk, wm, wn, aligned, batched, group_size, bits) \
   instantiate_kernel( \
-      #mode "_" #name "_" #type "_gs_32_b_4_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_alN_" #aligned "_batch_" #batched, \
+      #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_alN_" #aligned "_batch_" #batched, \
       fp_ ## name,    \
       type,    \
-      32,      \
-      4,       \
+      group_size,      \
+      bits,       \
       aligned, \
       batched)
 
-#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,      \
@@ -51,22 +51,27 @@
       transpose)
 
 
-#define instantiate_quantized_all_aligned(type) \
-  instantiate_quantized_aligned(mxfp4, gather_qmm_t_nax, type, 64, 64, 64, 2, 2, true)      \
-  instantiate_quantized_aligned(mxfp4, gather_qmm_t_nax, type, 64, 64, 64, 2, 2, false)     \
-  instantiate_quantized_aligned_batched(mxfp4, qmm_t_nax, type, 64, 64, 64, 2, 2, true, 1)  \
-  instantiate_quantized_aligned_batched(mxfp4, qmm_t_nax, type, 64, 64, 64, 2, 2, true, 0)  \
-  instantiate_quantized_aligned_batched(mxfp4, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 1) \
-  instantiate_quantized_aligned_batched(mxfp4, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 0)
+#define instantiate_quantized_all_aligned(type, mode, group_size, bits) \
+  instantiate_quantized_aligned(mode, gather_qmm_t_nax, type, 64, 64, 64, 2, 2, true, group_size, bits)      \
+  instantiate_quantized_aligned(mode, gather_qmm_t_nax, type, 64, 64, 64, 2, 2, false, group_size, bits)     \
+  instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, true, 1, group_size, bits)  \
+  instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, true, 0, group_size, bits)  \
+  instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 1, group_size, bits) \
+  instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 0, group_size, bits)
 
 
-#define instantiate_quantized_all_rhs(type) \
-  instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, mxfp4_gather_qmm_rhs_nax_nt, type, 64, 64, 64, 2, 2, true) \
-  instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, mxfp4_gather_qmm_rhs_nax_nn, type, 64, 64, 64, 2, 2, false) 
+#define instantiate_quantized_all_rhs(type, mode, group_size, bits) \
+  instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, gather_qmm_rhs_nax_nt, type, 64, 64, 64, 2, 2, true, mode, group_size, bits) \
+  instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, gather_qmm_rhs_nax_nn, type, 64, 64, 64, 2, 2, false, mode, group_size, bits)
+
+#define instantiate_quantized_modes(type, mode, group_size, bits) \
+  instantiate_quantized_all_aligned(type, mode, group_size, bits) \
+  instantiate_quantized_all_rhs(type, mode, group_size, bits)
 
 #define instantiate_quantized_types(type) \
-  instantiate_quantized_all_aligned(type) \
-  instantiate_quantized_all_rhs(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/indexing/masked_scatter.h

@@ -2,6 +2,8 @@
 
 #pragma once
 
+constant mlx::os_log logger("mlx", "masked_assign");
+
 template <typename T, bool src_contiguous>
 [[kernel]] void masked_assign_impl(
     const device bool* mask [[buffer(0)]],
@@ -21,6 +23,7 @@ template <typename T, bool src_contiguous>
 
   const uint src_index = scatter_offsets[idx];
   if (src_index >= src_batch_size) {
+    logger.log_debug("Out of bound read from src");
     return;
   }
 

mlx/backend/metal/kernels/logging.h

@@ -0,0 +1,26 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#if defined(__METAL_VERSION__) && (__METAL_VERSION__ >= 320)
+#include <metal_logging>
+
+namespace mlx {
+using os_log = metal::os_log;
+} // namespace mlx
+
+#else
+
+namespace mlx {
+struct os_log {
+  constexpr os_log(constant char*, constant char*) constant {}
+
+  template <typename... Args>
+  void log_debug(constant char*, Args...) const {}
+
+  template <typename... Args>
+  void log_debug(constant char*, Args...) const constant {}
+};
+} // namespace mlx
+
+#endif

mlx/backend/metal/kernels/utils.h

@@ -8,6 +8,7 @@
 #include "mlx/backend/metal/kernels/bf16_math.h"
 #include "mlx/backend/metal/kernels/complex.h"
 #include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/logging.h"
 
 typedef half float16_t;
 

mlx/backend/metal/quantized.cpp

@@ -315,9 +315,10 @@ void qvm_split_k(
   int B = out.size() / M / N;
   B *= split_k;
 
-  int bn = 64;
-  int bk = 32;
-  MTL::Size group_dims = MTL::Size(bk, 2, 1);
+  constexpr int num_simdgroups = 2;
+  constexpr int bk = 32;
+  int bn = std::min(group_size, 32) * num_simdgroups;
+  MTL::Size group_dims = MTL::Size(bk, num_simdgroups, 1);
   MTL::Size grid_dims = MTL::Size(M, N / bn, B);
 
   auto x_shape = x.shape();
@@ -431,9 +432,10 @@ void qvm(
     const std::string& mode) {
   int B = out.size() / M / N;
 
-  int bn = 64;
-  int bk = 32;
-  MTL::Size group_dims(bk, 2, 1);
+  constexpr int num_simdgroups = 2;
+  constexpr int bk = 32;
+  int bn = std::min(group_size, 32) * num_simdgroups;
+  MTL::Size group_dims(bk, num_simdgroups, 1);
   MTL::Size grid_dims(M, (N + bn - 1) / bn, B);
 
   std::string kname;
@@ -944,9 +946,10 @@ void gather_qvm(
     const std::string& mode) {
   int B = out.size() / M / N;
 
-  int bn = 64;
-  int bk = 32;
-  MTL::Size group_dims(bk, 2, 1);
+  constexpr int num_simdgroups = 2;
+  constexpr int bk = 32;
+  int bn = std::min(group_size, 32) * num_simdgroups;
+  MTL::Size group_dims(bk, num_simdgroups, 1);
   MTL::Size grid_dims(M, (N + bn - 1) / bn, B);
 
   std::string kname;
@@ -1452,25 +1455,26 @@ void fast::Quantize::eval_gpu(
   auto& compute_encoder = d.get_command_encoder(s.index);
 
   auto w = ensure_row_contiguous(w_pre, d, s);
-  compute_encoder.set_input_array(w, 0);
   if (dequantize_) {
     auto scales = ensure_row_contiguous(inputs[1], d, s);
-    compute_encoder.set_input_array(scales, 1);
-    compute_encoder.set_output_array(out, 3);
     if (mode_ == QuantizationMode::Affine) {
       auto biases = ensure_row_contiguous(inputs[2], d, s);
       compute_encoder.set_input_array(biases, 2);
     }
+    compute_encoder.set_input_array(w, 0);
+    compute_encoder.set_input_array(scales, 1);
+    compute_encoder.set_output_array(out, 3);
   } else {
     auto& scales = outputs[1];
     scales.set_data(allocator::malloc(scales.nbytes()));
-    compute_encoder.set_output_array(out, 1);
-    compute_encoder.set_output_array(scales, 2);
     if (mode_ == QuantizationMode::Affine) {
       auto& biases = outputs[2];
       biases.set_data(allocator::malloc(biases.nbytes()));
       compute_encoder.set_output_array(biases, 3);
     }
+    compute_encoder.set_input_array(w, 0);
+    compute_encoder.set_output_array(out, 1);
+    compute_encoder.set_output_array(scales, 2);
   }
 
   auto type_string = dequantize_ ? get_type_string(out.dtype())

mlx/io/load.cpp

@@ -72,6 +72,7 @@ Dtype dtype_from_array_protocol(std::string_view t) {
       case 'b': {
         if (size == 1)
           return bool_;
+        break;
       }
       case 'i': {
         if (size == 1)
@@ -82,6 +83,7 @@ Dtype dtype_from_array_protocol(std::string_view t) {
           return int32;
         else if (size == 8)
           return int64;
+        break;
       }
       case 'u': {
         if (size == 1)
@@ -92,21 +94,27 @@ Dtype dtype_from_array_protocol(std::string_view t) {
           return uint32;
         else if (size == 8)
           return uint64;
+        break;
       }
       case 'f': {
         if (size == 2)
           return float16;
         else if (size == 4)
           return float32;
+        else if (size == 8)
+          return float64;
+        break;
       }
       case 'c': {
-        return complex64;
+        if (size == 8)
+          return complex64;
+        break;
       }
     }
   }
 
   throw std::invalid_argument(
-      "[from_str] Invalid array protocol type-string: " + std::string(t));
+      "[from_str] Unsupported array protocol type-string: " + std::string(t));
 }
 
 #ifdef _WIN32

mlx/ops.cpp

@@ -4188,6 +4188,11 @@ std::pair<Dtype, QuantizationMode> validate_mode_with_type(
     } else {
       return {dtype, qmode};
     }
+  } else if (scales.dtype() != uint8) {
+    std::ostringstream msg;
+    msg << "[" << tag << "] Scale type must be uint8 but received type "
+        << scales.dtype() << ".";
+    throw std::invalid_argument(msg.str());
   }
   if (biases) {
     std::ostringstream msg;
@@ -6070,4 +6075,4 @@ array contiguous(
       {a});
 }
 
-} // namespace mlx::core
+} // namespace mlx::core

mlx/primitives.cpp

@@ -3397,9 +3397,11 @@ std::vector<array> QuantizedMatmul::vjp(
       throw std::runtime_error(
           "[QuantizedMatmul::vjp] no gradient wrt the quantized weights.");
     } else {
-      if (mode_ == QuantizationMode::Mxfp4) {
-        throw std::invalid_argument(
-            "[QuantizedMatmul::vjp] no gradient wrt scales with mxfp4 quantization.");
+      if (mode_ != QuantizationMode::Affine) {
+        std::ostringstream msg;
+        msg << "[QuantizedMatmul::vjp] no gradient wrt scales in "
+            << quantization_mode_to_string(mode_) << " quantization.";
+        throw std::invalid_argument(msg.str());
       }
       if (!dsb) {
         int ndim = primals[1].ndim();
@@ -3606,9 +3608,11 @@ std::vector<array> GatherQMM::vjp(
       throw std::runtime_error(
           "[GatherQMM::vjp] no gradient wrt the quantized weights.");
     } else {
-      if (mode_ == QuantizationMode::Mxfp4) {
-        throw std::invalid_argument(
-            "[GatherQMM::vjp] no gradient wrt scales with mxfp4 quantization.");
+      if (mode_ != QuantizationMode::Affine) {
+        std::ostringstream msg;
+        msg << "[GatherQMM::vjp] no gradient wrt scales in "
+            << quantization_mode_to_string(mode_) << " quantization.";
+        throw std::invalid_argument(msg.str());
       }
 
       if (!dsb) {

mlx/random.h

@@ -89,7 +89,7 @@ inline array uniform(
     const Shape& shape,
     const std::optional<array>& key = std::nullopt,
     StreamOrDevice s = {}) {
-  return uniform(shape, float32, key);
+  return uniform(shape, float32, key, s);
 }
 
 /** Generate samples from the standard normal distribution. */

mlx/version.h

@@ -4,7 +4,7 @@
 
 #define MLX_VERSION_MAJOR 0
 #define MLX_VERSION_MINOR 30
-#define MLX_VERSION_PATCH 1
+#define MLX_VERSION_PATCH 3
 #define MLX_VERSION_NUMERIC \
   (100000 * MLX_VERSION_MAJOR + 1000 * MLX_VERSION_MINOR + MLX_VERSION_PATCH)
 

pyproject.toml

@@ -1,7 +1,7 @@
 [build-system]
 requires = [
   "setuptools>=80",
-  "nanobind==2.10.2",
   "cmake>=3.25",
+  "typing_extensions",
 ]
 build-backend = "setuptools.build_meta"

python/mlx/_distributed_utils/config.py

@@ -17,6 +17,7 @@ from .common import (
     OptionalBoolAction,
     log,
     log_error,
+    log_warning,
     parse_hostfile,
     parse_hostlist,
 )
@@ -44,17 +45,29 @@ class ThunderboltHost:
     ports: list[ThunderboltPort]
 
 
-def add_ethernet_ips(hosts, verbose=False):
+def add_ips(hosts, verbose=False):
     # Get the ips for each host
     for h in hosts:
         log(verbose, "Getting the ip from", h.ssh_hostname)
-        h.ips.append(
-            run(
-                ["ssh", h.ssh_hostname, "ipconfig", "getifaddr", "en0"],
-                capture_output=True,
-                text=True,
-            ).stdout.strip()
-        )
+        ip = run(
+            ["ssh", h.ssh_hostname, "ipconfig", "getifaddr", "en0"],
+            capture_output=True,
+            text=True,
+        ).stdout.strip()
+        if ip != "":
+            h.ips.append(ip)
+            continue
+
+        ip = run(
+            ["ssh", h.ssh_hostname, "ipconfig", "getifaddr", "en1"],
+            capture_output=True,
+            text=True,
+        ).stdout.strip()
+        if ip != "":
+            h.ips.append(ip)
+            continue
+
+        log_warning("Could not extract ip for", h.ssh_hostname)
 
 
 def check_rdma(hosts, verbose=False):
@@ -393,7 +406,7 @@ def check_ssh_connections(hosts):
 
 def prepare_ethernet_hostfile(args, hosts):
     log(args.verbose, f"Preparing an ethernet hostfile")
-    add_ethernet_ips(hosts, args.verbose)
+    add_ips(hosts, args.verbose)
 
     hostfile = []
     for h in hosts:
@@ -438,7 +451,7 @@ def configure_ring(args, hosts, ips, ring, sshinfo):
 def configure_jaccl(args, hosts, ips, sshinfo):
     log(args.verbose, "Prepare a jaccl hostfile")
     check_rdma(hosts, args.verbose)
-    add_ethernet_ips(hosts, args.verbose)
+    add_ips(hosts, args.verbose)
 
     hostfile = []
     for i, h in enumerate(hosts):

python/mlx/_distributed_utils/launch.py

@@ -47,7 +47,7 @@ class CommandProcess:
 class RemoteProcess(CommandProcess):
     def __init__(self, rank, host, python, cwd, files, env, command):
         is_local = host == "127.0.0.1"
-        cmd = RemoteProcess.make_launch_script(rank, cwd, files, env, command)
+        cmd = RemoteProcess.make_launch_script(rank, cwd, files, env, command, is_local)
         if not is_local:
             cmd = f"ssh -tt -o LogLevel=QUIET {host} {shlex.quote(cmd)}"
 
@@ -104,11 +104,12 @@ class RemoteProcess(CommandProcess):
         self._killed = c.stdout.strip() == "1"
 
     @staticmethod
-    def make_launch_script(rank, cwd, files, env, command):
+    def make_launch_script(rank, cwd, files, env, command, is_local):
         script = ""
 
         # Disable echo
-        script = "stty -echo; "
+        if not is_local:
+            script = "stty -echo; "
 
         # Write the PID to a file so we can kill the process if needed
         script += "pidfile=$(mktemp); "

python/mlx/nn/init.py

@@ -100,7 +100,7 @@ def identity(dtype: mx.Dtype = mx.float32) -> Callable[[mx.array], mx.array]:
     r"""An initializer that returns an identity matrix.
 
     Args:
-        dtype (Dtype, optional): The data type of the array. Defaults:
+        dtype (Dtype, optional): The data type of the array. Default:
           ``float32``.
 
     Returns:
@@ -253,7 +253,7 @@ def he_normal(
     <https://arxiv.org/abs/1502.01852>`_
 
     Args:
-        dtype (Dtype, optional): The data type of the array. Defaults to mx.float32.
+        dtype (Dtype, optional): The data type of the array. Default: ``float32``.
 
     Returns:
         Callable[[array, str, float], array]: An initializer that returns an

python/mlx/nn/layers/__init__.py

@@ -87,7 +87,12 @@ from mlx.nn.layers.pooling import (
     MaxPool3d,
 )
 from mlx.nn.layers.positional_encoding import ALiBi, RoPE, SinusoidalPositionalEncoding
-from mlx.nn.layers.quantized import QuantizedEmbedding, QuantizedLinear, quantize
+from mlx.nn.layers.quantized import (
+    QQLinear,
+    QuantizedEmbedding,
+    QuantizedLinear,
+    quantize,
+)
 from mlx.nn.layers.recurrent import GRU, LSTM, RNN
 from mlx.nn.layers.transformer import (
     MultiHeadAttention,

python/mlx/nn/layers/activations.py

@@ -299,7 +299,7 @@ def hardswish(x):
 def hard_tanh(x, min_val=-1.0, max_val=1.0):
     r"""Applies the HardTanh function.
 
-    Applies :math:`\max(\min(x, \text{max\_val}), \text{min\_val})` element-wise.
+    Applies :math:`\max(\min(x, \mathrm{max\_val}), \mathrm{min\_val})` element-wise.
     """
     return mx.minimum(mx.maximum(x, min_val), max_val)
 

python/mlx/nn/layers/base.py

@@ -559,6 +559,9 @@ class Module(dict):
             _unfreeze_impl("", self)
         return self
 
+    def _set_training_mode(self, mode: bool) -> None:
+        self._training = mode
+
     def train(self, mode: bool = True) -> Module:
         """Set the model in or out of training mode.
 
@@ -573,10 +576,8 @@ class Module(dict):
             The module instance after updating the training mode.
         """
 
-        def _set_train(_, m):
-            m._training = mode
+        self.apply_to_modules(lambda _, m: m._set_training_mode(mode))
 
-        self.apply_to_modules(_set_train)
         return self
 
     def eval(self) -> Module:

python/mlx/nn/layers/embedding.py

@@ -1,6 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import math
+from typing import Optional
 
 import mlx.core as mx
 from mlx.nn.layers.base import Module
@@ -39,6 +40,11 @@ class Embedding(Module):
         """
         return x @ self.weight.T
 
-    def to_quantized(self, group_size: int = 64, bits: int = 4, mode: str = "affine"):
+    def to_quantized(
+        self,
+        group_size: Optional[int] = None,
+        bits: Optional[int] = None,
+        mode: str = "affine",
+    ):
         """Return a :obj:`QuantizedEmbedding` layer that approximates this embedding layer."""
         return QuantizedEmbedding.from_embedding(self, group_size, bits, mode)

python/mlx/nn/layers/linear.py

@@ -1,7 +1,7 @@
 # Copyright © 2023 Apple Inc.
 
 import math
-from typing import Any
+from typing import Any, Optional
 
 import mlx.core as mx
 from mlx.nn.layers.base import Module
@@ -70,7 +70,12 @@ class Linear(Module):
             x = x @ self["weight"].T
         return x
 
-    def to_quantized(self, group_size: int = 64, bits: int = 4, mode: str = "affine"):
+    def to_quantized(
+        self,
+        group_size: Optional[int] = None,
+        bits: Optional[int] = None,
+        mode: str = "affine",
+    ):
         """Return a :obj:`QuantizedLinear` layer that approximates this layer."""
         return QuantizedLinear.from_linear(self, group_size, bits, mode)
 

python/mlx/nn/layers/quantized.py

@@ -8,10 +8,21 @@ from mlx.nn.layers.base import Module
 from mlx.utils import tree_map_with_path
 
 
+def _defaults_for_mode(mode, group_size, bits):
+    mode_defaults = {
+        "affine": (64, 4),
+        "mxfp4": (32, 4),
+        "nvfp4": (16, 4),
+        "mxfp8": (32, 8),
+    }
+    default_group_size, default_bits = mode_defaults[mode]
+    return group_size or default_group_size, bits or default_bits
+
+
 def quantize(
     model: Module,
-    group_size: int = 64,
-    bits: int = 4,
+    group_size: int = None,
+    bits: int = None,
     *,
     mode: str = "affine",
     class_predicate: Optional[Callable[[str, Module], Union[bool, dict]]] = None,
@@ -24,10 +35,10 @@ def quantize(
 
     Args:
         model (mlx.nn.Module): The model whose leaf modules may be quantized.
-        group_size (int): The quantization group size (see
-           :func:`mlx.core.quantize`). Default: ``64``.
-        bits (int): The number of bits per parameter (see
-           :func:`mlx.core.quantize`). Default: ``4``.
+        group_size (Optional[int]): The quantization group size (see
+           :func:`mlx.core.quantize`). Default: ``None``.
+        bits (Optional[int]): The number of bits per parameter (see
+           :func:`mlx.core.quantize`). Default: ``None``.
         mode (str): The quantization method to use (see
            :func:`mlx.core.quantize`). Default: ``"affine"``.
         class_predicate (Optional[Callable]): A callable which receives the
@@ -72,10 +83,10 @@ class QuantizedEmbedding(Module):
         num_embeddings (int): How many possible discrete tokens can we embed.
            Usually called the vocabulary size.
         dims (int): The dimensionality of the embeddings.
-        group_size (int, optional): The group size to use for the quantized
-            weight. See :func:`~mlx.core.quantize`. Default: ``64``.
-        bits (int, optional): The bit width to use for the quantized weight.
-            See :func:`~mlx.core.quantize`. Default: ``4``.
+        group_size (Optional[int]): The group size to use for the quantized
+            weight. See :func:`~mlx.core.quantize`. Default: ``None``.
+        bits (Optional[int]): The bit width to use for the quantized weight.
+            See :func:`~mlx.core.quantize`. Default: ``None``.
         mode (str): The quantization method to use (see
            :func:`mlx.core.quantize`). Default: ``"affine"``.
     """
@@ -84,15 +95,14 @@ class QuantizedEmbedding(Module):
         self,
         num_embeddings: int,
         dims: int,
-        group_size: int = 64,
-        bits: int = 4,
+        group_size: int = None,
+        bits: int = None,
         mode: str = "affine",
     ):
         super().__init__()
 
         # Quantization config
-        self.group_size = group_size
-        self.bits = bits
+        self.group_size, self.bits = _defaults_for_mode(mode, group_size, bits)
         self.mode = mode
 
         # Initialize the quantized weight
@@ -147,8 +157,8 @@ class QuantizedEmbedding(Module):
     def from_embedding(
         cls,
         embedding_layer: Module,
-        group_size: int = 64,
-        bits: int = 4,
+        group_size: int = None,
+        bits: int = None,
         mode: str = "affine",
     ):
         """Create a :obj:`QuantizedEmbedding` layer from an :obj:`Embedding` layer."""
@@ -179,10 +189,10 @@ class QuantizedLinear(Module):
         output_dims (int): The dimensionality of the output features.
         bias (bool, optional): If set to ``False`` then the layer will not use
             a bias. Default: ``True``.
-        group_size (int, optional): The group size to use for the quantized
-            weight. See :func:`~mlx.core.quantize`. Default: ``64``.
-        bits (int, optional): The bit width to use for the quantized weight.
-            See :func:`~mlx.core.quantize`. Default: ``4``.
+        group_size (Optional[int]): The group size to use for the quantized
+            weight. See :func:`~mlx.core.quantize`. Default: ``None``.
+        bits (Optional[int]): The bit width to use for the quantized weight.
+            See :func:`~mlx.core.quantize`. Default: ``None``.
         mode (str): The quantization method to use (see
            :func:`mlx.core.quantize`). Default: ``"affine"``.
     """
@@ -192,15 +202,14 @@ class QuantizedLinear(Module):
         input_dims: int,
         output_dims: int,
         bias: bool = True,
-        group_size: int = 64,
-        bits: int = 4,
+        group_size: int = None,
+        bits: int = None,
         mode: str = "affine",
     ):
         super().__init__()
 
         # Quantization config
-        self.group_size = group_size
-        self.bits = bits
+        self.group_size, self.bits = _defaults_for_mode(mode, group_size, bits)
         self.mode = mode
 
         # Initialize the quantized weight
@@ -249,8 +258,8 @@ class QuantizedLinear(Module):
     def from_linear(
         cls,
         linear_layer: Module,
-        group_size: int = 64,
-        bits: int = 4,
+        group_size: int = None,
+        bits: int = None,
         mode: str = "affine",
     ):
         """Create a :obj:`QuantizedLinear` layer from a :obj:`Linear` layer."""
@@ -268,3 +277,128 @@ class QuantizedLinear(Module):
             ql.bias = linear_layer.bias
 
         return ql
+
+
+class QQLinear(Module):
+    """Quantizes the input and applies an affine transformation using quantized weights.
+
+    Two use cases are supported:
+
+    1) **Eval**:  The weights are frozen and stored in quantized form together with
+       their scales (``self.weight`` is quantized and ``self.scales`` is provided).
+    2) **Train**: The weights are stored in higher precision and are quantized on
+         the fly during computation so that gradients with respect to the weights
+         can be computed.
+
+    To switch between the two cases, use ``layer.eval()`` and ``layer.train()`` respectively.
+
+    Compared to the :class:`mlx.nn.QuantizedLinear` layer, this layer
+    quantizes the input as well and includes weights in gradient computations.
+
+    :obj:`QQLinear` also provides:
+     -  the class method :meth:`from_linear` to convert :class:`mlx.nn.Linear`
+     layers to :obj:`QQLinear` layers.
+
+    Note: This layer does not support a bias term yet.
+
+    Args:
+        input_dims (int): The dimensionality of the input features.
+        output_dims (int): The dimensionality of the output features.
+        group_size (Optional[int]): The group size to use for the quantized weight.
+            See :func:`~mlx.core.quantize`. Default: ``None``.
+        bits (Optional[int]): The bit width to use for the quantized weight.
+            See :func:`~mlx.core.quantize`. Default: ``None``.
+        mode (Optional[str]): The quantization method to use (see
+            :func:`mlx.core.quantize`). Currently, only ``"nvfp4"`` and ``"mxfp8"``
+            are supported. Default: ``"nvfp4"``.
+    """
+
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: int,
+        group_size: int = None,
+        bits: int = None,
+        mode: str = "nvfp4",
+    ):
+        super().__init__()
+
+        # Quantization config
+        self.group_size, self.bits = _defaults_for_mode(mode, group_size, bits)
+        self.mode = mode
+
+        scale = math.sqrt(1 / input_dims)
+        self.weight = mx.random.uniform(
+            low=-scale,
+            high=scale,
+            shape=(output_dims, input_dims),
+        )
+        self._quantized = False
+
+    def _extra_repr(self):
+        out_dims, in_dims = self.weight.shape
+        if self.weight.dtype == mx.uint32:
+            in_dims *= 32 // self.bits
+        return (
+            f"input_dims={in_dims}, output_dims={out_dims}, "
+            f"group_size={self.group_size}, bits={self.bits}, mode={self.mode}"
+        )
+
+    def quantize(self):
+        if not self._quantized:
+            self.weight, self.scales = mx.quantize(
+                self.weight,
+                self.group_size,
+                self.bits,
+                mode=self.mode,
+            )
+            self._quantized = True
+
+    def dequantize(self):
+        if self._quantized:
+            self.weight = mx.dequantize(
+                self.weight,
+                scales=self.scales,
+                group_size=self.group_size,
+                bits=self.bits,
+                mode=self.mode,
+            )
+            self.__delattr__("scales")
+            self._quantized = False
+
+    def _set_training_mode(self, mode: bool):
+        super()._set_training_mode(mode)
+
+        if self._training:
+            self.dequantize()
+        else:
+            self.quantize()
+
+    def __call__(self, x):
+        x = mx.qqmm(
+            x,
+            self["weight"],
+            scales=self.get("scales"),
+            group_size=self.group_size,
+            bits=self.bits,
+            mode=self.mode,
+        )
+        return x
+
+    @classmethod
+    def from_linear(
+        cls,
+        linear_layer: Module,
+        group_size: int = None,
+        bits: int = None,
+        mode: str = "nvfp4",
+    ):
+        """Create a :obj:`QQLinear` layer from a :obj:`Linear` layer."""
+        output_dims, input_dims = linear_layer.weight.shape  # (N,K)
+        if linear_layer.get("bias") is not None:
+            raise NotImplementedError("QQLinear does not support bias yet.")
+        ql = cls(input_dims, output_dims, group_size, bits, mode=mode)
+        ql.weight = linear_layer.weight
+        ql.train(linear_layer.training)
+
+        return ql

python/scripts/repair_cuda.sh

@@ -1,26 +0,0 @@
-#!/bin/bash
-
-auditwheel repair dist/* \
-  --plat manylinux_2_35_${1} \
-  --exclude libcublas* \
-  --exclude libnvrtc* \
-  --exclude libcuda* \
-  --exclude libcudnn* \
-  --exclude libnccl* \
-  -w wheel_tmp
-
-
-mkdir wheelhouse
-cd wheel_tmp
-repaired_wheel=$(find . -name "*.whl" -print -quit)
-unzip -q "${repaired_wheel}"
-rm "${repaired_wheel}"
-mlx_so="mlx/lib/libmlx.so"
-rpath=$(patchelf --print-rpath "${mlx_so}")
-base="\$ORIGIN/../../nvidia"
-rpath=$rpath:${base}/cublas/lib:${base}/cuda_nvrtc/lib:${base}/cudnn/lib:${base}/nccl/lib
-patchelf --force-rpath --set-rpath "$rpath" "$mlx_so"
-python ../python/scripts/repair_record.py ${mlx_so}
-
-# Re-zip the repaired wheel
-zip -r -q "../wheelhouse/${repaired_wheel}" .

python/scripts/repair_linux.sh

@@ -1,20 +0,0 @@
-#!/bin/bash
-
-auditwheel repair dist/* \
-  --plat manylinux_2_35_${1} \
-  --only-plat \
-  --exclude libmlx* \
-  -w wheel_tmp
-
-mkdir wheelhouse
-cd wheel_tmp
-repaired_wheel=$(find . -name "*.whl" -print -quit)
-unzip -q "${repaired_wheel}"
-rm "${repaired_wheel}"
-core_so=$(find mlx -name "core*.so" -print -quit)
-rpath="\$ORIGIN/lib"
-patchelf --force-rpath --set-rpath "$rpath" "$core_so"
-python ../python/scripts/repair_record.py ${core_so}
-
-# Re-zip the repaired wheel
-zip -r -q "../wheelhouse/${repaired_wheel}" .

python/scripts/repair_record.py

@@ -1,33 +0,0 @@
-import base64
-import glob
-import hashlib
-import sys
-
-filename = sys.argv[1]
-
-
-# Compute the new hash and size
-def urlsafe_b64encode(data: bytes) -> bytes:
-    return base64.urlsafe_b64encode(data).rstrip(b"=")
-
-
-hasher = hashlib.sha256()
-with open(filename, "rb") as f:
-    data = f.read()
-    hasher.update(data)
-hash_str = urlsafe_b64encode(hasher.digest()).decode("ascii")
-size = len(data)
-
-# Update the record file
-record_file = glob.glob("*/RECORD")[0]
-with open(record_file, "r") as f:
-    lines = [l.split(",") for l in f.readlines()]
-
-for l in lines:
-    if filename == l[0]:
-        l[1] = hash_str
-        l[2] = f"{size}\n"
-
-with open(record_file, "w") as f:
-    for l in lines:
-        f.write(",".join(l))

python/src/CMakeLists.txt

@@ -29,6 +29,37 @@ nanobind_add_module(
   ${CMAKE_CURRENT_SOURCE_DIR}/trees.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp)
 
+if(MLX_BUILD_PYTHON_STUBS)
+  nanobind_add_stub(
+    core_stub
+    # Run stubgen -m mlx.core -i python -p _stub_patterns.txt -o python/mlx
+    RECURSIVE
+    MODULE
+    "mlx.core"
+    PYTHON_PATH
+    "${CMAKE_CURRENT_SOURCE_DIR}/.."
+    PATTERN_FILE
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/_stub_patterns.txt"
+    OUTPUT_PATH
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx"
+    # Note that the list is passed to cmake for dependency managment and not
+    # used by stubgen.
+    OUTPUT
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/__init__.pyi"
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/cuda.pyi"
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/distributed.pyi"
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/fast.pyi"
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/fft.pyi"
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/linalg.pyi"
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/metal.pyi"
+    "${CMAKE_CURRENT_SOURCE_DIR}/../mlx/core/random.pyi"
+    # Make this an optional installable component.
+    EXCLUDE_FROM_ALL
+    INSTALL_TIME
+    COMPONENT
+    core_stub)
+endif()
+
 if(NOT MLX_PYTHON_BINDINGS_OUTPUT_DIRECTORY)
   if(NOT CMAKE_LIBRARY_OUTPUT_DIRECTORY)
     set(MLX_PYTHON_BINDINGS_OUTPUT_DIRECTORY ${PROJECT_BINARY_DIR})
@@ -55,8 +86,10 @@ target_link_libraries(core PRIVATE mlx)
 
 if(BUILD_SHARED_LIBS)
   if(${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
-    target_link_options(core PRIVATE -Wl,-rpath,@loader_path/lib)
+    set_target_properties(core PROPERTIES INSTALL_RPATH "@loader_path/lib")
   else()
-    target_link_options(core PRIVATE -Wl,-rpath,\$ORIGIN/lib)
+    set_target_properties(core PROPERTIES INSTALL_RPATH "\$ORIGIN/lib")
   endif()
+  # Do not add build dir to rpath.
+  set_target_properties(core PROPERTIES BUILD_WITH_INSTALL_RPATH ON)
 endif()

python/tests/cuda_skip.py

@@ -6,7 +6,7 @@ cuda_skip = {
     # Gather matmul NYI
     "TestBlas.test_gather_matmul",
     "TestBlas.test_gather_matmul_grad",
-    "TestBlas.test_gather_mm_sorted",
+    "TestBlas.test_gather_mm_sorted_vjp",
     # Segmented matmul NYI
     "TestBlas.test_segmented_mm",
     # Hadamard NYI
@@ -48,8 +48,8 @@ cuda_skip = {
     "TestQuantized.test_qmm_shapes",
     "TestQuantized.test_qmm_vjp",
     "TestQuantized.test_qmv",
-    "TestQuantized.test_mxfp4_qmv",
-    "TestQuantized.test_mxfp4_qvm",
+    "TestQuantized.test_fp_qmv",
+    "TestQuantized.test_fp_qvm",
     "TestQuantized.test_qvm",
     "TestQuantized.test_qvm_splitk",
     "TestQuantized.test_small_matrix",

python/tests/test_blas.py

@@ -1235,15 +1235,39 @@ class TestBlas(mlx_tests.MLXTestCase):
         def gather_mm_test(a, b, rhs):
             return mx.gather_mm(a, b, rhs_indices=rhs, sorted_indices=True)
 
+        dtypes = [(mx.float32, 1e-4)]
+        if mx.cuda.is_available():
+            dtypes += [
+                (mx.float16, 1e-3),
+                (mx.bfloat16, 1e-2),
+            ]
+
+        for b_transposed in (True, False):
+            for dtype, tol in dtypes:
+                with self.subTest(b_transposed=b_transposed, dtype=dtype):
+                    a = mx.random.normal((100, 1, 100), dtype=dtype)
+                    b = mx.random.normal((8, 100, 100), dtype=dtype)
+                    if b_transposed:
+                        b = b.swapaxes(-1, -2)
+                    rhs = mx.sort(mx.random.randint(0, 8, shape=(100,)))
+
+                    c1 = gather_mm_ref(a, b, rhs)
+                    c2 = gather_mm_test(a, b, rhs)
+                    self.assertTrue(mx.allclose(c1, c2, rtol=tol, atol=tol))
+
+    def test_gather_mm_sorted_vjp(self):
+        def gather_mm_ref(a, b, rhs):
+            b = b[rhs]
+            return a @ b
+
+        def gather_mm_test(a, b, rhs):
+            return mx.gather_mm(a, b, rhs_indices=rhs, sorted_indices=True)
+
         a = mx.random.normal((100, 1, 100))
         b = mx.random.normal((8, 100, 100))
         rhs = mx.sort(mx.random.randint(0, 8, shape=(100,)))
 
-        c1 = gather_mm_ref(a, b, rhs)
-        c2 = gather_mm_test(a, b, rhs)
-        self.assertTrue(mx.allclose(c1, c2, atol=1e-4))
-
-        cotan = mx.random.normal(c1.shape)
+        cotan = mx.random.normal((100, 1, 100))
         c1, dc1 = mx.vjp(
             lambda a, b: gather_mm_ref(a, b, rhs),
             [a, b],

python/tests/test_load.py

@@ -71,6 +71,21 @@ class TestLoad(mlx_tests.MLXTestCase):
         load_arr = mx.load(Path(save_file))
         self.assertTrue(mx.array_equal(load_arr, save_arr))
 
+    def test_load_npy_dtype(self):
+        save_file = os.path.join(self.test_dir, "mlx_path.npy")
+        a = np.random.randn(8).astype(np.float64)
+        np.save(save_file, a)
+        out = mx.load(save_file, stream=mx.cpu)
+        self.assertEqual(out.dtype, mx.float64)
+        self.assertTrue(np.array_equal(np.array(out), a))
+
+        a = np.random.randn(8).astype(np.float64)
+        b = np.random.randn(8).astype(np.float64)
+        c = a + 0j * b
+        np.save(save_file, c)
+        with self.assertRaises(Exception):
+            out = mx.load(save_file, stream=mx.cpu)
+
     def test_save_and_load_safetensors(self):
         test_file = os.path.join(self.test_dir, "test.safetensors")
         with self.assertRaises(Exception):

python/tests/test_quantized.py

@@ -289,7 +289,7 @@ class TestQuantized(mlx_tests.MLXTestCase):
             [128, 64, 32],  # group_size
             [2, 3, 4, 5, 6, 8],  # bits
             [256, 512, 67],  # M
-            [64, 128],  # N
+            [64, 256],  # N
             [0, 1, 3, 8],  # B
         )
         for group_size, bits, M, N, B in tests:
@@ -309,33 +309,34 @@ class TestQuantized(mlx_tests.MLXTestCase):
                 self.assertEqual(y_q.shape, y_hat.shape)
                 self.assertLess((y_q - y_hat).abs().max(), 1e-3)
 
-    def test_mxfp4_qmv(self):
+    def test_fp_qmv(self):
         key = mx.random.key(0)
         k1, k2 = mx.random.split(key)
         tests = product(
             [256, 512, 67],  # M
-            [64, 128],  # N
+            [64, 256],  # N
             [0, 1, 3, 8],  # B
         )
+        modes = ["mxfp4", "nvfp4", "mxfp8"]
         for M, N, B in tests:
-            with self.subTest(shape=(B, M, N), group_size=32):
-                x_shape = (3, 1, N) if B == 0 else (B, 1, N)
-                w_shape = (M, N) if B == 0 else (B, M, N)
-                x = mx.random.normal(shape=x_shape, key=k1)
-                w = mx.random.normal(shape=w_shape, key=k2)
-                w_q, scales = mx.quantize(w, group_size=32, mode="mxfp4")
-                w_hat = mx.dequantize(w_q, scales, group_size=32, mode="mxfp4")
-                y_q = mx.quantized_matmul(
-                    x,
-                    w_q,
-                    scales,
-                    transpose=True,
-                    group_size=32,
-                    mode="mxfp4",
-                )
-                y_hat = x @ mx.swapaxes(w_hat, -1, -2)
-                self.assertEqual(y_q.shape, y_hat.shape)
-                self.assertLess((y_q - y_hat).abs().max(), 1e-3)
+            for mode in modes:
+                with self.subTest(shape=(B, M, N), mode=mode):
+                    x_shape = (3, 1, N) if B == 0 else (B, 1, N)
+                    w_shape = (M, N) if B == 0 else (B, M, N)
+                    x = mx.random.normal(shape=x_shape, key=k1)
+                    w = mx.random.normal(shape=w_shape, key=k2)
+                    w_q, scales = mx.quantize(w, mode=mode)
+                    w_hat = mx.dequantize(w_q, scales, mode=mode)
+                    y_q = mx.quantized_matmul(
+                        x,
+                        w_q,
+                        scales,
+                        transpose=True,
+                        mode=mode,
+                    )
+                    y_hat = x @ mx.swapaxes(w_hat, -1, -2)
+                    self.assertEqual(y_q.shape, y_hat.shape)
+                    self.assertLess((y_q - y_hat).abs().max(), 1e-3)
 
     def test_qvm(self):
         key = mx.random.key(0)
@@ -402,7 +403,7 @@ class TestQuantized(mlx_tests.MLXTestCase):
         self.assertEqual(y_q.shape, y_hat.shape)
         self.assertLess((y_q - y_hat).abs().max(), 2e-3)
 
-    def test_mxfp4_qvm(self):
+    def test_fp_qvm(self):
         key = mx.random.key(0)
         k1, k2 = mx.random.split(key)
         tests = product(
@@ -413,26 +414,27 @@ class TestQuantized(mlx_tests.MLXTestCase):
         # Add a splitk
         tests = list(tests)
         tests.append((128, 16384, 0))
+        modes = ["mxfp4", "nvfp4", "mxfp8"]
 
         for M, N, B in tests:
-            with self.subTest(shape=(B, M, N)):
-                x_shape = (1, N) if B == 0 else (B, 1, N)
-                w_shape = (N, M) if B == 0 else (B, N, M)
-                x = mx.random.normal(shape=x_shape, key=k1)
-                w = mx.random.normal(shape=w_shape, key=k2)
-                w_q, scales = mx.quantize(w, group_size=32, mode="mxfp4")
-                w_hat = mx.dequantize(w_q, scales, group_size=32, mode="mxfp4")
-                y_q = mx.quantized_matmul(
-                    x,
-                    w_q,
-                    scales,
-                    transpose=False,
-                    group_size=32,
-                    mode="mxfp4",
-                )
-                y_hat = x @ w_hat
-                self.assertEqual(y_q.shape, y_hat.shape)
-                self.assertLess((y_q - y_hat).abs().max(), 2e-3)
+            for mode in modes:
+                with self.subTest(shape=(B, M, N), mode=mode):
+                    x_shape = (1, N) if B == 0 else (B, 1, N)
+                    w_shape = (N, M) if B == 0 else (B, N, M)
+                    x = mx.random.normal(shape=x_shape, key=k1)
+                    w = mx.random.normal(shape=w_shape, key=k2)
+                    w_q, scales = mx.quantize(w, mode=mode)
+                    w_hat = mx.dequantize(w_q, scales, mode=mode)
+                    y_q = mx.quantized_matmul(
+                        x,
+                        w_q,
+                        scales,
+                        transpose=False,
+                        mode=mode,
+                    )
+                    y_hat = x @ w_hat
+                    self.assertEqual(y_q.shape, y_hat.shape)
+                    self.assertLess((y_q - y_hat).abs().max(), 2e-3)
 
     def test_mode_error_cases(self):
         w = mx.random.normal(shape=(256, 256))
@@ -626,7 +628,7 @@ class TestQuantized(mlx_tests.MLXTestCase):
         self.assertLess((y_q - y_hat).abs().max(), 1e-3)
 
     def test_gather_qmm(self):
-        def quantize(w, transpose=True, group_size=64, bits=4, mode="affine"):
+        def quantize(w, transpose=True, group_size=None, bits=None, mode="affine"):
             if mode == "affine":
                 qw, s, b = mx.quantize(w, group_size=group_size, bits=bits, mode=mode)
             else:
@@ -647,8 +649,8 @@ class TestQuantized(mlx_tests.MLXTestCase):
             lhs_indices=None,
             rhs_indices=None,
             transpose=True,
-            group_size=64,
-            bits=4,
+            group_size=None,
+            bits=None,
             mode="affine",
         ):
             with self.subTest(
@@ -737,9 +739,22 @@ class TestQuantized(mlx_tests.MLXTestCase):
                 "lhs_indices": (0,),
                 "batch_B": (3,),
                 "rhs_indices": (2, 1),
-                "group_size": 32,
+                "mode": "nvfp4",
+            },
+            {
+                "batch_A": (1,),
+                "lhs_indices": (0,),
+                "batch_B": (3,),
+                "rhs_indices": (2, 1),
                 "mode": "mxfp4",
             },
+            {
+                "batch_A": (1,),
+                "lhs_indices": (0,),
+                "batch_B": (3,),
+                "rhs_indices": (2, 1),
+                "mode": "mxfp8",
+            },
         )
 
         for kwargs in inputs:
@@ -753,24 +768,24 @@ class TestQuantized(mlx_tests.MLXTestCase):
             test_shape(32, 512, 32, transpose=False, **kwargs)
             test_shape(1, 512, 32, transpose=False, **kwargs)
 
-    def test_qmm_mxfp4_type(self):
+    def test_qmm_fp_type(self):
         indices = mx.array([[2], [0], [1]], dtype=mx.uint32)
 
-        for t in [mx.bfloat16, mx.float16, mx.float32]:
-            x = mx.random.normal((32, 256)).astype(t)
+        modes = ["mxfp8", "mxfp4"]
+        for mode in modes:
+            for t in [mx.bfloat16, mx.float16, mx.float32]:
+                x = mx.random.normal((32, 256)).astype(t)
 
-            w = mx.random.normal((32, 256))
-            wq, s = mx.quantize(w, mode="mxfp4", bits=4, group_size=32)
-            out = mx.quantized_matmul(x, wq, s, mode="mxfp4", group_size=32, bits=4)
-            self.assertEqual(out.dtype, t)
+                w = mx.random.normal((32, 256))
+                wq, s = mx.quantize(w, mode=mode)
+                out = mx.quantized_matmul(x, wq, s, mode=mode)
+                self.assertEqual(out.dtype, t)
 
-            w = mx.random.normal((4, 32, 256))
-            wq, s = mx.quantize(w, mode="mxfp4", bits=4, group_size=32)
+                w = mx.random.normal((4, 32, 256))
+                wq, s = mx.quantize(w, mode=mode)
 
-            out = mx.gather_qmm(
-                x, wq, s, rhs_indices=indices, mode="mxfp4", group_size=32, bits=4
-            )
-            self.assertEqual(out.dtype, t)
+                out = mx.gather_qmm(x, wq, s, rhs_indices=indices, mode=mode)
+                self.assertEqual(out.dtype, t)
 
     def test_gather_matmul_grad(self):
         def quantize(w, transpose=True, group_size=64, bits=4):
@@ -802,14 +817,14 @@ class TestQuantized(mlx_tests.MLXTestCase):
         self.assertTrue(mx.allclose(g1, g2, atol=1e-4))
 
     def test_gather_qmm_sorted(self):
-        def quantize(w, transpose=True, bits=4, group_size=64, mode="affine"):
+        def quantize(w, transpose=True, group_size=None, mode="affine"):
             if mode == "affine":
-                qw, s, b = mx.quantize(w, group_size=group_size, bits=bits, mode=mode)
+                qw, s, b = mx.quantize(w, group_size=group_size, mode=mode)
             else:
-                qw, s = mx.quantize(w, group_size=group_size, bits=bits, mode=mode)
+                qw, s = mx.quantize(w, mode=mode)
                 b = None
 
-            w_hat = mx.dequantize(qw, s, b, group_size=group_size, bits=bits, mode=mode)
+            w_hat = mx.dequantize(qw, s, b, group_size=group_size, mode=mode)
             if transpose:
                 w_hat = w_hat.swapaxes(-1, -2)
             return w_hat, qw, s, b
@@ -831,11 +846,15 @@ class TestQuantized(mlx_tests.MLXTestCase):
             # L, K, D, E, I, transpose
             (32, 512, 512, 4, 2, True, "affine"),
             (32, 512, 544, 4, 2, True, "mxfp4"),
+            (32, 512, 544, 4, 2, True, "nvfp4"),
+            (32, 512, 544, 4, 2, True, "mxfp8"),
             (133, 512, 512, 4, 2, True, "affine"),
             (133, 512, 555, 4, 2, True, "affine"),
             (133, 512, 512, 4, 2, True, "affine"),
             (64, 512, 512, 4, 2, False, "affine"),
             (64, 512, 544, 4, 2, False, "mxfp4"),
+            (64, 512, 544, 4, 2, False, "nvfp4"),
+            (64, 512, 544, 4, 2, False, "mxfp8"),
             (133, 512, 512, 4, 2, False, "affine"),
             (133, 512, 544, 4, 2, False, "affine"),
             (133, 512, 555, 4, 2, False, "affine"),
@@ -848,8 +867,8 @@ class TestQuantized(mlx_tests.MLXTestCase):
 
         for L, K, D, E, I, transpose, mode in parameters:
             with self.subTest(L=L, K=K, D=D, E=E, I=I, transpose=transpose, mode=mode):
-                if mode == "mxfp4":
-                    group_size = 32
+                if mode != "affine":
+                    group_size = None
                     dtype = (
                         mx.bfloat16 if (mx.default_device() == mx.gpu) else mx.float32
                     )
@@ -984,36 +1003,44 @@ class TestQuantized(mlx_tests.MLXTestCase):
                 num_ds = (out_up - out_down) / (2 * eps)
                 self.assertAlmostEqual(dparams[p][idx], num_ds, delta=2e-2)
 
-    def test_mxfp4_vjp_scales_throws(self):
+    def test_fp_vjp_scales_throws(self):
         mx.random.seed(0)
         x = mx.random.normal(shape=(2, 512))
         w = mx.random.normal(shape=(512, 512))
-        wq, s = mx.quantize(w, bits=4, group_size=32, mode="mxfp4")
+        for mode in ["mxfp4", "mxfp8", "nvfp4"]:
+            wq, s = mx.quantize(w, mode=mode)
 
-        def mm(s, x, wq):
-            return mx.quantized_matmul(
-                x, wq, s, bits=4, group_size=32, mode="mxfp4"
-            ).sum()
+            def mm(s, x, wq):
+                return mx.quantized_matmul(x, wq, s, mode=mode).sum()
 
-        # Should raise
-        with self.assertRaises(ValueError):
-            ds = mx.grad(mm)(s, x, wq)
+            # Should raise
+            with self.assertRaises(ValueError):
+                ds = mx.grad(mm)(s, x, wq)
 
-        rhs_indices = mx.array(0)
-        with self.assertRaises(ValueError):
+            rhs_indices = mx.array(0)
+            with self.assertRaises(ValueError):
 
-            def gmm(s, x, wq):
-                return mx.gather_qmm(
-                    x,
-                    wq,
-                    s,
-                    rhs_indices=rhs_indices,
-                    bits=4,
-                    group_size=32,
-                    mode="mxfp4",
-                ).sum()
+                def gmm(s, x, wq):
+                    return mx.gather_qmm(
+                        x,
+                        wq,
+                        s,
+                        rhs_indices=rhs_indices,
+                        mode=mode,
+                    ).sum()
 
-            ds = mx.grad(gmm)(s, x, wq)
+                ds = mx.grad(gmm)(s, x, wq)
+
+    def test_quantize_strided(self):
+        N = 64
+        mode = "nvfp4"
+        w = mx.random.normal(shape=(N, N))
+        w_q, scales = mx.quantize(w, mode="nvfp4")
+
+        scales = mx.broadcast_to(mx.array(56, mx.uint8), scales.shape)
+        w_hat = mx.dequantize(w_q, scales, mode=mode)
+        expected = mx.dequantize(w_q, mx.contiguous(scales), mode=mode)
+        self.assertTrue(mx.allclose(w_hat, expected))
 
 
 if __name__ == "__main__":

setup.py

@@ -8,7 +8,7 @@ import subprocess
 from functools import partial
 from pathlib import Path
 
-from setuptools import Command, Extension, find_namespace_packages, setup
+from setuptools import Extension, find_namespace_packages, setup
 from setuptools.command.bdist_wheel import bdist_wheel
 from setuptools.command.build_ext import build_ext
 
@@ -79,22 +79,22 @@ class CMakeBuild(build_ext):
         if not build_temp.exists():
             build_temp.mkdir(parents=True)
 
-        build_python = "ON"
-        install_prefix = f"{extdir}{os.sep}"
+        install_prefix = extdir
+        pybind_out_dir = extdir
         if build_stage == 1:
             # Don't include MLX libraries in the wheel
-            install_prefix = f"{build_temp}"
+            install_prefix = build_temp
         elif build_stage == 2:
             # Don't include Python bindings in the wheel
-            build_python = "OFF"
+            pybind_out_dir = build_temp
         cmake_args = [
             f"-DCMAKE_INSTALL_PREFIX={install_prefix}",
+            f"-DMLX_PYTHON_BINDINGS_OUTPUT_DIRECTORY={pybind_out_dir}",
             f"-DCMAKE_BUILD_TYPE={cfg}",
-            f"-DMLX_BUILD_PYTHON_BINDINGS={build_python}",
+            "-DMLX_BUILD_PYTHON_BINDINGS=ON",
             "-DMLX_BUILD_TESTS=OFF",
             "-DMLX_BUILD_BENCHMARKS=OFF",
             "-DMLX_BUILD_EXAMPLES=OFF",
-            f"-DMLX_PYTHON_BINDINGS_OUTPUT_DIRECTORY={extdir}{os.sep}",
         ]
         if build_stage == 2 and build_cuda:
             # Last arch is always real and virtual for forward-compatibility
@@ -155,46 +155,25 @@ class CMakeBuild(build_ext):
     def run(self):
         super().run()
 
+        ext = next(ext for ext in self.extensions if ext.name == "mlx.core")
+
         # Based on https://github.com/pypa/setuptools/blob/main/setuptools/command/build_ext.py#L102
         if self.inplace:
-            for ext in self.extensions:
-                if ext.name == "mlx.core":
-                    # Resolve inplace package dir
-                    build_py = self.get_finalized_command("build_py")
-                    inplace_file, regular_file = self._get_inplace_equivalent(
-                        build_py, ext
-                    )
+            # Resolve inplace package dir
+            build_py = self.get_finalized_command("build_py")
+            inplace_file, regular_file = self._get_inplace_equivalent(build_py, ext)
 
-                    inplace_dir = str(Path(inplace_file).parent.resolve())
-                    regular_dir = str(Path(regular_file).parent.resolve())
+            inplace_dir = str(Path(inplace_file).parent.resolve())
+            regular_dir = str(Path(regular_file).parent.resolve())
 
-                    self.copy_tree(regular_dir, inplace_dir)
+            self.copy_tree(regular_dir, inplace_dir)
 
-
-class GenerateStubs(Command):
-    user_options = []
-
-    def initialize_options(self):
-        pass
-
-    def finalize_options(self):
-        pass
-
-    def run(self) -> None:
-        out_path = "python/mlx/core"
-        stub_cmd = [
-            "python",
-            "-m",
-            "nanobind.stubgen",
-            "-m",
-            "mlx.core",
-            "-p",
-            "python/mlx/_stub_patterns.txt",
-        ]
-        subprocess.run(stub_cmd + ["-r", "-O", out_path])
-        # Run again without recursive to specify output file name
-        subprocess.run(["rm", f"{out_path}/mlx.pyi"])
-        subprocess.run(stub_cmd + ["-o", f"{out_path}/__init__.pyi"])
+        # Build type stubs.
+        build_temp = Path(self.build_temp) / ext.name
+        subprocess.run(
+            ["cmake", "--install", build_temp, "--component", "core_stub"],
+            check=True,
+        )
 
 
 class MLXBdistWheel(bdist_wheel):
@@ -246,7 +225,6 @@ if __name__ == "__main__":
         ext_modules=[CMakeExtension("mlx.core")],
         cmdclass={
             "build_ext": CMakeBuild,
-            "generate_stubs": GenerateStubs,
             "bdist_wheel": MLXBdistWheel,
         },
     )
@@ -255,11 +233,9 @@ if __name__ == "__main__":
 
     extras = {
         "dev": [
-            "nanobind==2.10.2",
-            "numpy",
+            "numpy>=2",
             "pre-commit",
-            "setuptools>=80",
-            "torch",
+            "torch>=2.9",
             "typing_extensions",
         ],
     }
@@ -313,6 +289,9 @@ if __name__ == "__main__":
         elif build_cuda:
             toolkit = cuda_toolkit_major_version()
             name = f"mlx-cuda-{toolkit}"
+            # Note: update following files when new dependency is added:
+            # * .github/actions/build-cuda-release/action.yml
+            # * mlx/backend/cuda/CMakeLists.txt
             if toolkit == 12:
                 install_requires += [
                     "nvidia-cublas-cu12==12.9.*",
@@ -320,8 +299,8 @@ if __name__ == "__main__":
                 ]
             elif toolkit == 13:
                 install_requires += [
-                    "nvidia-cublas-cu13",
-                    "nvidia-cuda-nvrtc-cu13",
+                    "nvidia-cublas",
+                    "nvidia-cuda-nvrtc",
                 ]
             else:
                 raise ValueError(f"Unknown toolkit {toolkit}")

tests/CMakeLists.txt

@@ -37,5 +37,7 @@ target_sources(
           ${METAL_TEST_SOURCES})
 
 target_link_libraries(tests PRIVATE mlx doctest)
+target_compile_options(tests PRIVATE ${SANITIZER_COMPILE_FLAGS})
+target_link_options(tests PRIVATE ${SANITIZER_LINK_FLAGS})
 doctest_discover_tests(tests)
 add_test(NAME tests COMMAND tests)

tests/array_tests.cpp

@@ -613,7 +613,12 @@ TEST_CASE("test make array from user buffer") {
   std::vector<int> buffer(size, 0);
 
   int count = 0;
-  auto deleter = [&count](void*) { count++; };
+  auto deleter = [&count, data = buffer.data()](void* ptr) {
+    // make sure pointer is correct
+    if (ptr == data) {
+      count++;
+    }
+  };
 
   {
     auto a = array(buffer.data(), Shape{size}, int32, deleter);