[CUDA] Use qmv kernel for fp quantizations (#3239 )

Add initial tuning for M5 pro and max (#3211 )
[CUDA] Implement SegmentedMM (#3238 )
2026-03-12 07:25:17 +09:00 · 2026-03-11 14:05:43 -07:00 · 2026-03-11 13:31:43 -07:00 · 2026-03-11 19:45:51 +09:00 · 2026-03-11 12:58:04 +09:00 · 2026-03-10 11:47:25 +01:00
321 changed files with 15290 additions and 4693 deletions
@@ -20,7 +20,7 @@ runs:
      run: |
        pip install auditwheel build patchelf setuptools
        python setup.py clean --all
-        MLX_BUILD_STAGE=2 python -m build -w
+        MLX_DISABLE_SM90A_KERNELS=1 MLX_BUILD_STAGE=2 python -m build -w

        auditwheel repair dist/mlx_cuda*.whl \
          --plat manylinux_2_35_${{ inputs.arch }} \
@@ -9,6 +9,7 @@ inputs:
 runs:
  using: "composite"
  steps:
+
    - name: Install Python package
      id: python_build
      shell: sh
@@ -20,10 +21,12 @@ runs:
      run: |
        if ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }} ; then
          # There is no GPU in arm64 runner, use a common arch.
-          CMAKE_ARGS="$CMAKE_ARGS -DMLX_CUDA_ARCHITECTURES=90a"
+          CMAKE_ARGS="$CMAKE_ARGS -DMLX_CUDA_ARCHITECTURES=80"
          # Can not build tests and stubs when the built executables can not run.
          CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF -DMLX_BUILD_PYTHON_STUBS=OFF"
        fi
+        # Install cpu-only torch to save space
+        pip install torch --index-url https://download.pytorch.org/whl/cpu
        pip install --no-build-isolation -e ".[dev]" -v
        # Pass the CMAKE_ARGS to following steps.
        echo CMAKE_ARGS="$CMAKE_ARGS" >> $GITHUB_OUTPUT
@@ -18,6 +18,7 @@ runs:
    - name: Build Python package
      shell: bash -l {0}
      env:
+        DEVELOPER_DIR: /Applications/Xcode-latest.app
        MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
      run: |
        pip install build
@@ -28,6 +29,7 @@ runs:
      if: ${{ inputs.build-backend }}
      shell: bash -l {0}
      env:
+        DEVELOPER_DIR: /Applications/Xcode-latest.app
        MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
      run: |
        python setup.py clean --all
@@ -17,15 +17,15 @@ runs:
    - name: Install tests dependencies
      shell: bash -l {0}
      run: |
-        pip install numpy torch tensorflow unittest-xml-reporting
+        pip install numpy torch tensorflow

    - name: Run Python tests
      shell: bash -l {0}
      env:
        LOW_MEMORY: 1
      run: |
-        DEVICE=cpu python -m xmlrunner discover -v python/tests -o test-results/cpu
-        DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -m xmlrunner discover -v python/tests -o test-results/gpu
+        DEVICE=cpu python -m unittest discover -v python/tests
+        DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -m unittest discover -v python/tests
        mpirun --bind-to none -host localhost:8 -np 8 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python python/tests/mpi_test_distributed.py
        mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
        if $(grep "\[WARN\]" stderr.log); then echo "Distributed ring test failed"; exit 1; fi
@@ -77,6 +77,4 @@ runs:
      run: |
        CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
          pip install -e . -v
-        python -m xmlrunner discover \
-            -v python/tests \
-            -o test-results/gpu_jit
+        python -m unittest discover -v python/tests
@@ -0,0 +1,26 @@
+name: 'Build on Windows'
+
+runs:
+  using: 'composite'
+  steps:
+    - name: Install Python package
+      id: python-build
+      shell: cmd
+      env:
+        # For MSVC, Ninja/Release is the only config supported by ccache.
+        CMAKE_ARGS: >-
+          -G Ninja
+          -DCMAKE_BUILD_TYPE=Release
+          -DCMAKE_C_COMPILER=cl
+          -DCMAKE_CXX_COMPILER=cl
+          -DCMAKE_RC_COMPILER=rc
+      run: |
+        uv pip install ".[dev]" -v
+        :: Pass the CMAKE_ARGS to following steps.
+        >>%GITHUB_OUTPUT% ECHO CMAKE_ARGS=%CMAKE_ARGS%
+
+    - name: Build CPP only
+      shell: cmd
+      run: |
+        cmake . -B build ${{ steps.python-build.outputs.CMAKE_ARGS }}
+        cmake --build build -j %NUMBER_OF_PROCESSORS%
@@ -0,0 +1,42 @@
+name: 'Setup Windows environment'
+
+inputs:
+  python-version:
+    description: 'Version of python to set up'
+    required: false
+    default: '3.14'
+  use-ccache:
+    description: 'Whether to enable ccache'
+    required: false
+    default: 'true'
+
+runs:
+  using: 'composite'
+  steps:
+    - name: Use ccache
+      if: ${{ inputs.use-ccache == 'true' }}
+      uses: hendrikmuhs/ccache-action@v1.2
+      with:
+        key: ccache-${{ runner.os }}-${{ runner.arch }}-cpu
+        max-size: 1GB
+
+    - name: Setup Visual Studio cmd
+      shell: cmd
+      run: |
+        :: Find out path to VS.
+        pushd "C:\Program Files (x86)\Microsoft Visual Studio\Installer\"
+        for /f "delims=" %%x in ('.\vswhere.exe -latest -property InstallationPath') do set VSPATH=%%x
+        popd
+        :: Import VS vars.
+        call "%VSPATH%\VC\Auxiliary\Build\vcvarsall.bat" x64
+        :: Export to all steps.
+        >>%GITHUB_ENV% set
+
+    - uses: astral-sh/setup-uv@v7
+
+    - name: Setup Python venv
+      shell: cmd
+      run: |
+        uv venv --python ${{ inputs.python-version }}
+        call ".venv/Scripts/activate.bat"
+        >>%GITHUB_ENV% set
@@ -0,0 +1,20 @@
+name: 'Run tests on Windows'
+
+runs:
+  using: 'composite'
+  steps:
+    - name: Run Python tests - CPU
+      shell: bash
+      run: |
+        echo "::group::Python tests - CPU"
+        python -m unittest discover python/tests -v
+        echo "::endgroup::"
+
+    - name: Run CPP tests - CPU
+      shell: bash
+      env:
+        DEVICE: cpu
+      run: |
+        echo "::group::CPP tests - CPU"
+        ./build/tests.exe -tce="*gguf*,test random uniform"
+        echo "::endgroup::"
@@ -36,6 +36,7 @@ jobs:
      - uses: ./.github/actions/setup-linux
      - uses: ./.github/actions/build-linux
      - uses: ./.github/actions/test-linux
+      - run: df -h

  cuda_build_and_test:
    name: Linux (${{ matrix.toolkit }}, ${{ matrix.arch }})
@@ -75,6 +76,16 @@ jobs:
      - uses: ./.github/actions/setup-macos
      - uses: ./.github/actions/build-macos

+  windows_build_and_test:
+    name: Windows (cpu, x86_64)
+    needs: check_lint
+    runs-on: windows-2025
+    steps:
+      - uses: actions/checkout@v6
+      - uses: ./.github/actions/setup-windows
+      - uses: ./.github/actions/build-windows
+      - uses: ./.github/actions/test-windows
+
  build_documentation:
    name: Build Documentation
    if: github.repository == 'ml-explore/mlx'
@@ -23,18 +23,19 @@ jobs:
          build-backend: ${{ matrix.python-version == '3.10' }}
          arch: "x86_64"
      - name: Upload mlx artifacts
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
        with:
          name: linux-wheels-${{ matrix.python_version }}
          path: wheelhouse/mlx-*.whl
          retention-days: 7
      - name: Upload mlx-cpu artifacts
        if: matrix.python_version == '3.10'
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
        with:
          name: mlx-cpu
          path: wheelhouse/mlx_cpu-*.whl
          retention-days: 7
+      - run: df -h

  build_linux_with_tests:
    strategy:
@@ -52,6 +53,7 @@ jobs:
          python-version: ${{ matrix.python_version }}
      - uses: ./.github/actions/build-linux
      - uses: ./.github/actions/test-linux
+      - run: df -h

  build_mac_release:
    if: github.repository == 'ml-explore/mlx'
@@ -95,7 +97,7 @@ jobs:
          toolkit: 'cuda-12.9'
          arch: 'x86_64'
      - name: Upload artifacts
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
        with:
          name: mlx-cuda
          path: wheelhouse/mlx_cuda_*.whl
@@ -4,18 +4,18 @@ on:
  push:
    tags:
      - 'v*'
+    branches:
+      - 'test-publish/*'
  workflow_dispatch:
    inputs:
-      publish:
-        description: 'Publish to PyPI (uncheck for dry run)'
+      dry_run:
+        description: 'Dry run (do not publish to PyPi)'
        required: false
        type: boolean
-        default: true
      dev_release:
        description: 'Development release (DEV_RELEASE=1)'
        required: false
        type: boolean
-        default: false

 permissions:
  contents: read
@@ -29,7 +29,7 @@ jobs:
      - uses: ./.github/actions/build-docs

  deploy_documentation:
-    if: inputs.publish
+    if: ${{ !inputs.dry_run }}
    needs: build_documentation
    permissions:
      pages: write
@@ -64,7 +64,7 @@ jobs:
          build-backend: ${{ matrix.python_version == '3.10' }}
          arch: ${{ matrix.arch }}
      - name: Upload MLX artifacts
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
        with:
          overwrite: true
          name: linux-wheels-${{ matrix.python_version }}-${{ matrix.arch }}
@@ -72,7 +72,7 @@ jobs:
          if-no-files-found: error
      - name: Upload CPU artifacts
        if: matrix.python_version == '3.10'
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
        with:
          overwrite: true
          name: mlx-cpu-${{ matrix.arch }}
@@ -110,8 +110,13 @@ jobs:
        with:
          macos-target: 15.0
          build-backend: ${{ matrix.python-version == '3.10' }}
+      - name: Build macOS 26 package
+        uses: ./.github/actions/build-macos-release
+        with:
+          macos-target: 26.0
+          build-backend: ${{ matrix.python-version == '3.10' }}
      - name: Upload MLX artifacts
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
        with:
          overwrite: true
          name: mac-wheels-${{ matrix.python-version }}
@@ -119,7 +124,7 @@ jobs:
          if-no-files-found: error
      - name: Upload Metal artifacts
        if: matrix.python-version == '3.10'
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
        with:
          overwrite: true
          name: mlx-metal
@@ -147,7 +152,7 @@ jobs:
        with:
          arch: ${{ matrix.arch }}
      - name: Upload artifacts
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
        with:
          overwrite: true
          name: mlx-${{ matrix.toolkit }}-${{ matrix.arch }}
@@ -161,15 +166,15 @@ jobs:
    permissions:
      id-token: write
    environment:
-      name: ${{ inputs.publish && 'pypi' || '' }}
+      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
      url: https://pypi.org/p/mlx
    steps:
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
        with:
          pattern: linux-wheels-*
          merge-multiple: true
          path: dist
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
        with:
          pattern: mac-wheels-*
          merge-multiple: true
@@ -177,7 +182,7 @@ jobs:
      - name: Display structure of downloaded files
        run: du -ah dist
      - name: Publish package distributions to PyPI
-        if: inputs.publish
+        if: ${{ !inputs.dry_run }}
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
          repository-url: https://upload.pypi.org/legacy/
@@ -189,10 +194,10 @@ jobs:
    permissions:
      id-token: write
    environment:
-      name: ${{ inputs.publish && 'pypi' || '' }}
+      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
      url: https://pypi.org/p/mlx-cuda
    steps:
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
        with:
          pattern: mlx-cuda-*
          merge-multiple: true
@@ -200,7 +205,7 @@ jobs:
      - name: Display structure of downloaded files
        run: du -ah dist
      - name: Publish package distributions to PyPI
-        if: inputs.publish
+        if: ${{ !inputs.dry_run }}
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
          repository-url: https://upload.pypi.org/legacy/
@@ -212,10 +217,10 @@ jobs:
    permissions:
      id-token: write
    environment:
-      name: ${{ inputs.publish && 'pypi' || '' }}
+      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
      url: https://pypi.org/p/mlx-cpu
    steps:
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
        with:
          pattern: mlx-cpu-*
          merge-multiple: true
@@ -223,7 +228,7 @@ jobs:
      - name: Display structure of downloaded files
        run: du -ah dist
      - name: Publish package distributions to PyPI
-        if: inputs.publish
+        if: ${{ !inputs.dry_run }}
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
          repository-url: https://upload.pypi.org/legacy/
@@ -235,17 +240,17 @@ jobs:
    permissions:
      id-token: write
    environment:
-      name: ${{ inputs.publish && 'pypi' || '' }}
+      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
      url: https://pypi.org/p/mlx-metal
    steps:
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
        with:
          name: mlx-metal
          path: dist
      - name: Display structure of downloaded files
        run: du -ah dist
      - name: Publish package distributions to PyPI
-        if: inputs.publish
+        if: ${{ !inputs.dry_run }}
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
          repository-url: https://upload.pypi.org/legacy/
@@ -3,16 +3,12 @@ __pycache__/
 *.py[cod]
 *$py.class

-# C extensions
-*.so
-
 # tensor files
 *.safe
 *.safetensors

 # Metal libraries
 *.metallib
-venv/

 # Distribution / packaging
 python/mlx/core
@@ -30,6 +26,7 @@ lib64/
 parts/
 sdist/
 var/
+venv/
 wheels/
 share/python-wheels/
 *.egg-info/
@@ -37,12 +34,7 @@ share/python-wheels/
 *.egg
 MANIFEST
 uv.lock
-
-# vim
-*.swp
-
-# Ignore build dir
-build/
+.DS_Store

 # Prerequisites
 *.d
@@ -52,6 +44,7 @@ build/
 *.lo
 *.o
 *.obj
+*.ilk

 # Precompiled Headers
 *.gch
@@ -80,9 +73,9 @@ build/
 # Debug symbols
 *.pdb

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

 # ----------------------------- Setup -----------------------------
 set(CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
-set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD 20)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 set(CMAKE_INSTALL_MESSAGE NEVER)
@@ -40,7 +40,6 @@ option(MLX_METAL_DEBUG "Enhance metal debug workflow" OFF)
 option(MLX_ENABLE_X64_MAC "Enable building for x64 macOS" OFF)
 option(MLX_BUILD_GGUF "Include support for GGUF format" ON)
 option(MLX_BUILD_SAFETENSORS "Include support for safetensors format" ON)
-option(MLX_BUILD_BLAS_FROM_SOURCE "Build OpenBLAS from source code" OFF)
 option(MLX_BUILD_PYTHON_STUBS "Build stub files for python bindings" ON)
 option(MLX_METAL_JIT "Use JIT compilation for Metal kernels" OFF)
 option(MLX_USE_CCACHE "Use CCache for compilation cache when available" ON)
@@ -150,15 +149,13 @@ cmake_policy(SET CMP0135 NEW)

 add_library(mlx)

-# Supress warnings: note: parameter passing for argument of type
-# ‘std::pair<float, float>’ when C++17 is enabled changed to match C++14 in GCC
-# 10.1
-target_compile_options(mlx PRIVATE -Wno-psabi)
 target_compile_options(mlx PUBLIC ${SANITIZER_COMPILE_FLAGS})
 target_link_options(mlx PUBLIC ${SANITIZER_LINK_FLAGS})

 if(MLX_BUILD_CUDA)
  enable_language(CUDA)
+  find_package(CUDAToolkit REQUIRED)
+  find_package(CUDNN REQUIRED)
 endif()

 if(MLX_BUILD_METAL)
@@ -222,14 +219,17 @@ if(WIN32)
  if(MSVC)
    # GGUF does not build with MSVC.
    set(MLX_BUILD_GGUF OFF)
-    # There is no prebuilt OpenBLAS distribution for MSVC.
-    set(MLX_BUILD_BLAS_FROM_SOURCE ON)
+  endif()
+  # Generate DLL and EXE in the same dir, otherwise EXE will not be able to run.
+  # This is only done when MLX is built as the top project.
+  if(CMAKE_CURRENT_SOURCE_DIR STREQUAL CMAKE_SOURCE_DIR)
+    set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR})
  endif()
  # Windows implementation of dlfcn.h APIs.
  FetchContent_Declare(
    dlfcn-win32
    GIT_REPOSITORY https://github.com/dlfcn-win32/dlfcn-win32.git
-    GIT_TAG v1.4.1
+    GIT_TAG v1.4.2
    EXCLUDE_FROM_ALL)
  block()
  set(BUILD_SHARED_LIBS OFF)
@@ -253,20 +253,25 @@ if(MLX_BUILD_CPU)
    target_link_libraries(mlx PUBLIC ${ACCELERATE_LIBRARY})
    add_compile_definitions(MLX_USE_ACCELERATE)
    add_compile_definitions(ACCELERATE_NEW_LAPACK)
-  elseif(MLX_BUILD_BLAS_FROM_SOURCE)
-    # Download and build OpenBLAS from source code.
+  elseif(WIN32)
+    # Download and link prebuilt binaries of OpenBLAS. Note that we can only
+    # link with the dynamic library, the prebuilt binaries were built with MinGW
+    # so static-linking would require linking with MinGW's runtime.
    FetchContent_Declare(
      openblas
-      GIT_REPOSITORY https://github.com/OpenMathLib/OpenBLAS.git
-      GIT_TAG v0.3.28
-      EXCLUDE_FROM_ALL)
-    set(BUILD_STATIC_LIBS ON) # link statically
-    set(NOFORTRAN ON) # msvc has no fortran compiler
+      URL "https://github.com/OpenMathLib/OpenBLAS/releases/download/v0.3.31/OpenBLAS-0.3.31-x64.zip"
+    )
    FetchContent_MakeAvailable(openblas)
-    target_link_libraries(mlx PRIVATE openblas)
-    target_include_directories(
-      mlx PRIVATE "${openblas_SOURCE_DIR}/lapack-netlib/LAPACKE/include"
-                  "${CMAKE_BINARY_DIR}/generated" "${CMAKE_BINARY_DIR}")
+    target_link_libraries(mlx
+                          PRIVATE "${openblas_SOURCE_DIR}/lib/libopenblas.lib")
+    target_include_directories(mlx PRIVATE "${openblas_SOURCE_DIR}/include")
+    # Make sure the DLL file is placed in the same dir with executables.
+    set(OPENBLAS_DLL_FILE "${openblas_SOURCE_DIR}/bin/libopenblas.dll")
+    add_custom_command(
+      TARGET mlx
+      POST_BUILD
+      COMMAND ${CMAKE_COMMAND} -E copy_if_different ${OPENBLAS_DLL_FILE}
+              ${CMAKE_BINARY_DIR})
  else()
    if(${CMAKE_HOST_APPLE})
      # The blas shipped in macOS SDK is not supported, search homebrew for
@@ -318,16 +323,13 @@ target_include_directories(
  mlx PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_LIST_DIR}>
             $<INSTALL_INTERFACE:include>)

-# Do not add mlx_EXPORTS define for shared library.
-set_target_properties(mlx PROPERTIES DEFINE_SYMBOL "")
-
 if(USE_SYSTEM_FMT)
  find_package(fmt REQUIRED)
 else()
  FetchContent_Declare(
    fmt
    GIT_REPOSITORY https://github.com/fmtlib/fmt.git
-    GIT_TAG 10.2.1
+    GIT_TAG 12.1.0
    EXCLUDE_FROM_ALL)
  FetchContent_MakeAvailable(fmt)
 endif()
@@ -365,6 +367,15 @@ endif()
 # ----------------------------- Installation -----------------------------
 include(GNUInstallDirs)

+if(WIN32)
+  # Install DLLs to the same dir with extension file (core.pyd) on Windows.
+  set(CMAKE_INSTALL_BINDIR ".")
+  if(MLX_BUILD_CPU)
+    # Install OpenBLAS.
+    install(FILES ${OPENBLAS_DLL_FILE} TYPE BIN)
+  endif()
+endif()
+
 # Install library
 install(
  TARGETS mlx
@@ -0,0 +1,152 @@
+import math
+import time
+
+import mlx.core as mx
+import numpy as np
+import torch
+
+N_warmup = 2
+N_iter_bench = 10
+N_iter_func = 10
+
+
+def bench(f, a, b, b_prime):
+    for i in range(N_warmup):
+        f(a, b, b_prime)
+    torch.mps.synchronize()
+
+    s = time.perf_counter_ns()
+    for i in range(N_iter_bench):
+        f(a, b, b_prime)
+    e = time.perf_counter_ns()
+    return (e - s) * 1e-9
+
+
+def make_mx_conv_3D(strides=(1, 1, 1), padding=(0, 0, 0), groups=1):
+    def mx_conv_3D(a, b, b_prime):
+        y = a
+        for i in range(N_iter_func):
+            y = mx.conv3d(y, b, stride=strides, padding=padding, groups=groups)
+            y = mx.conv3d(y, b_prime, stride=strides, padding=padding, groups=groups)
+        mx.eval(y)
+        return y
+
+    return mx_conv_3D
+
+
+def make_pt_conv_3D(strides=(1, 1, 1), padding=(0, 0, 0), groups=1):
+    @torch.no_grad()
+    def pt_conv_3D(a, b, b_prime):
+        y = a
+        for i in range(N_iter_func):
+            y = torch.conv3d(y, b, stride=strides, padding=padding, groups=groups)
+            y = torch.conv3d(y, b_prime, stride=strides, padding=padding, groups=groups)
+        torch.mps.synchronize()
+        return y
+
+    return pt_conv_3D
+
+
+def bench_shape(N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype):
+    scale = 1.0 / math.sqrt(kD * kH * kW * C)
+    a_np = np.random.uniform(0, 0.5, (N, D, H, W, C))
+    b_np = np.random.uniform(-scale, scale, (O, kD, kH, kW, int(C / groups)))
+    b_prime_np = np.random.uniform(-scale, scale, (C, kD, kH, kW, int(O / groups)))
+
+    a_np, b_np, b_prime_np = map(lambda x: x.astype(np_dtype), (a_np, b_np, b_prime_np))
+    a_mx, b_mx, b_prime_mx = map(lambda x: mx.array(x), (a_np, b_np, b_prime_np))
+    a_pt, b_pt, b_prime_pt = map(
+        lambda x: torch.from_numpy(x.transpose(0, 4, 1, 2, 3)).to("mps"),
+        (a_np, b_np, b_prime_np),
+    )
+
+    torch.mps.synchronize()
+
+    f_mx = make_mx_conv_3D(strides, padding, groups)
+    f_pt = make_pt_conv_3D(strides, padding, groups)
+
+    time_torch = bench(f_pt, a_pt, b_pt, b_prime_pt)
+    time_mlx = bench(f_mx, a_mx, b_mx, b_prime_mx)
+
+    # Measure MLX memory
+    mx.clear_cache()
+    mx.reset_peak_memory()
+    y = mx.conv3d(a_mx, b_mx, stride=strides, padding=padding, groups=groups)
+    mx.eval(y)
+    mlx_peak_mb = mx.get_peak_memory() / 1024**2
+    mlx_active_mb = mx.get_active_memory() / 1024**2
+    del y
+
+    # Measure PyTorch MPS memory
+    torch.mps.synchronize()
+    torch.mps.empty_cache()
+    y = torch.conv3d(a_pt, b_pt, stride=strides, padding=padding, groups=groups)
+    torch.mps.synchronize()
+    pt_current_mb = torch.mps.current_allocated_memory() / 1024**2
+    pt_driver_mb = torch.mps.driver_allocated_memory() / 1024**2
+    del y
+
+    out_mx = mx.conv3d(a_mx, b_mx, stride=strides, padding=padding, groups=groups)
+    out_pt = torch.conv3d(
+        a_pt.to("cpu"), b_pt.to("cpu"), stride=strides, padding=padding, groups=groups
+    )
+    out_pt = torch.permute(out_pt, (0, 2, 3, 4, 1))
+    out_pt = out_pt.numpy(force=True)
+
+    atol = 2e-5 if np_dtype == np.float32 else 5e-4
+
+    if not np.allclose(out_pt, out_mx, atol=atol):
+        print(
+            f"Failed at {(N, D, H, W, C)}, {(O, kD, kH, kW, C)} "
+            f"[strides = {strides}, padding = {padding}, groups = {groups}] "
+            f"with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}"
+        )
+
+    return time_mlx, time_torch, mlx_peak_mb, mlx_active_mb, pt_current_mb, pt_driver_mb
+
+
+if __name__ == "__main__":
+    dtypes = ("float16", "float32")
+    shapes = (
+        # (C % 16 == 0)
+        (4, 16, 16, 16, 32, 3, 3, 3, 32, (1, 1, 1), (1, 1, 1), 1),
+        (4, 16, 16, 16, 64, 3, 3, 3, 64, (1, 1, 1), (1, 1, 1), 1),
+        (4, 16, 16, 16, 128, 3, 3, 3, 128, (1, 1, 1), (1, 1, 1), 1),
+        (4, 32, 32, 32, 64, 3, 3, 3, 64, (1, 1, 1), (1, 1, 1), 1),
+        (4, 32, 32, 32, 128, 3, 3, 3, 128, (1, 1, 1), (1, 1, 1), 1),
+        # Larger spatial dims
+        (2, 64, 64, 64, 32, 3, 3, 3, 64, (1, 1, 1), (1, 1, 1), 1),
+        (1, 64, 64, 64, 64, 3, 3, 3, 128, (1, 1, 1), (1, 1, 1), 1),
+        # Strided
+        (4, 32, 32, 32, 64, 3, 3, 3, 128, (2, 2, 2), (1, 1, 1), 1),
+        # Asymmetric kernels
+        (4, 32, 32, 32, 64, 3, 1, 1, 128, (1, 1, 1), (1, 0, 0), 1),
+        (4, 32, 32, 32, 64, 1, 3, 3, 128, (1, 1, 1), (0, 1, 1), 1),
+        # (C % 16 != 0)
+        (4, 16, 16, 16, 21, 3, 3, 3, 21, (1, 1, 1), (1, 1, 1), 1),
+        (4, 16, 16, 16, 55, 3, 3, 3, 55, (1, 1, 1), (1, 1, 1), 1),
+        (4, 32, 32, 32, 55, 3, 3, 3, 55, (1, 1, 1), (1, 1, 1), 1),
+        (4, 16, 16, 16, 3, 3, 3, 3, 32, (1, 1, 1), (1, 1, 1), 1),
+    )
+
+    for dtype in dtypes:
+        print(f"\n{'=' * 120}" f"\n  dtype: {dtype}" f"\n{'=' * 120}")
+        print(
+            f"{'(N,   D,   H,   W,   C)':<26s} {'(  O, kD, kH, kW,   C)':<24s} "
+            f"{'stride':<12s} {'pads':<12s} {'groups':>6s} "
+            f"{'diff%':>7s}  "
+            f"{'MLX peak':>9s} {'MLX act':>8s} {'PT cur':>8s} {'PT drv':>8s}"
+        )
+        for N, D, H, W, C, kD, kH, kW, O, strides, padding, groups in shapes:
+            np_dtype = getattr(np, dtype)
+            time_mlx, time_torch, mlx_peak, mlx_act, pt_cur, pt_drv = bench_shape(
+                N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype
+            )
+            diff = time_torch / time_mlx - 1.0
+
+            print(
+                f"({N}, {D:3d}, {H:3d}, {W:3d}, {C:3d}), ({O:3d}, {kD:2d}, {kH:2d}, {kW:2d}, {C:3d}), "
+                f"{strides}, {padding}, {groups:6d}, "
+                f"{100. * diff:+6.1f}%  "
+                f"{mlx_peak:8.1f}  {mlx_act:7.1f}  {pt_cur:7.1f}  {pt_drv:7.1f}"
+            )
@@ -0,0 +1,119 @@
+# Copyright © 2026 Apple Inc.
+
+import math
+import time
+
+import mlx.core as mx
+import numpy as np
+import torch
+
+N_WARMUP = 5
+N_BENCH = 20
+
+
+def bench_mlx(a, b):
+    for _ in range(N_WARMUP):
+        mx.eval(a @ b)
+
+    times = []
+    for _ in range(N_BENCH):
+        start = time.perf_counter_ns()
+        mx.eval(a @ b)
+        end = time.perf_counter_ns()
+        times.append((end - start) * 1e-9)
+
+    return np.mean(times), np.std(times)
+
+
+@torch.no_grad()
+def bench_torch(a, b):
+    for _ in range(N_WARMUP):
+        _ = a @ b
+        torch.mps.synchronize()
+
+    times = []
+    for _ in range(N_BENCH):
+        start = time.perf_counter_ns()
+        _ = a @ b
+        torch.mps.synchronize()
+        end = time.perf_counter_ns()
+        times.append((end - start) * 1e-9)
+
+    return np.mean(times), np.std(times)
+
+
+def check_correctness(out_mx, out_pt, rtol, M, N, K):
+    if not np.allclose(out_pt, out_mx, rtol=rtol, atol=0):
+        abs_diff = np.abs(out_pt - out_mx)
+        rel_diff = abs_diff / np.maximum(np.abs(out_pt), 1e-10)
+
+        print(
+            f"  WARNING: Correctness failed at {M}x{N}x{K}: "
+            f"max_abs={np.max(abs_diff):.6e}, max_rel={np.max(rel_diff):.6e}"
+        )
+
+
+def bench_gemm(M, N, K, dtype, rtol):
+    scale = 0.5 / math.sqrt(K)
+    a_np = np.random.uniform(0, scale, (M, K)).astype(np.float32)
+    b_np = np.random.uniform(0, scale, (K, N)).astype(np.float32)
+
+    a_mx = mx.array(a_np).astype(getattr(mx, dtype))
+    b_mx = mx.array(b_np).astype(getattr(mx, dtype))
+
+    a_pt = torch.from_numpy(a_np).to(dtype=getattr(torch, dtype), device="mps")
+    b_pt = torch.from_numpy(b_np).to(dtype=getattr(torch, dtype), device="mps")
+    torch.mps.synchronize()
+
+    torch_mean, torch_std = bench_torch(a_pt, b_pt)
+    mlx_mean, mlx_std = bench_mlx(a_mx, b_mx)
+
+    out_mx = (a_mx @ b_mx).astype(mx.float32)
+    out_pt = (a_pt @ b_pt).to(torch.float32).to("cpu").numpy(force=True)
+    check_correctness(out_mx, out_pt, rtol, M, N, K)
+
+    return mlx_mean, mlx_std, torch_mean, torch_std
+
+
+if __name__ == "__main__":
+    dtypes = ("bfloat16", "float16", "float32")
+
+    rtols = {
+        "float32": 1e-3,
+        "float16": 5e-3,
+        "bfloat16": 1e-2,
+    }
+
+    shapes = (
+        (2048, 2048, 10240),
+        (2048, 3072, 10240),
+        (3072, 3072, 10240),
+        (3072, 3072, 12288),
+        (3072, 4096, 12288),
+        (4096, 4096, 12288),
+        (4096, 4096, 18432),
+        (4096, 4096, 21504),
+        (4096, 6144, 21504),
+        (6144, 6144, 21504),
+    )
+
+    for dtype in dtypes:
+        print(f"\nPerformance ({dtype}):")
+        print(
+            f"{'M':>5s} {'N':>5s} {'K':>6s}  "
+            f"{'MLX (ms)':>15s}  {'Torch (ms)':>15s}  {'Speedup':>10s}"
+        )
+        print("-" * 80)
+
+        for M, N, K in shapes:
+            mlx_mean, mlx_std, torch_mean, torch_std = bench_gemm(
+                M, N, K, dtype, rtols[dtype]
+            )
+            speedup = torch_mean / mlx_mean
+
+            print(
+                f"{M:5d} {N:5d} {K:6d}  "
+                f"{mlx_mean*1000:7.2f}±{mlx_std*1000:5.2f}  "
+                f"{torch_mean*1000:7.2f}±{torch_std*1000:5.2f}  "
+                f"{speedup:8.2f}x"
+            )
@@ -0,0 +1,209 @@
+# Copyright © 2026 Apple Inc.
+
+import argparse
+import time
+
+import mlx.core as mx
+import numpy as np
+
+MLX_DTYPES = {
+    "float16": mx.float16,
+    "bfloat16": mx.bfloat16,
+    "float32": mx.float32,
+}
+
+
+def parse_cases(cases):
+    parsed = []
+    for spec in cases.split(","):
+        m, n, k, s = [int(x) for x in spec.split("x")]
+        parsed.append((m, n, k, s))
+    return parsed
+
+
+def make_segments(k, num_segments, pattern, seed):
+    if pattern == "equal":
+        cuts = np.linspace(0, k, num_segments + 1, dtype=np.int64)
+    else:
+        rng = np.random.default_rng(seed)
+        cuts = rng.integers(0, k + 1, size=(num_segments - 1,), dtype=np.int64)
+        cuts = np.sort(cuts)
+        cuts = np.concatenate(([0], cuts, [k]))
+    return np.stack([cuts[:-1], cuts[1:]], axis=1).astype(np.uint32)
+
+
+def numpy_segmented_mm_ref(a, b, segments):
+    """Ground-truth reference in float64."""
+    out = []
+    for start, end in segments:
+        out.append(a[:, start:end] @ b[start:end, :])
+    return np.stack(out, axis=0)
+
+
+def mlx_segmented_mm_loop(a, b, segments):
+    """MLX loop-of-matmuls baseline."""
+    segments_list = segments.tolist()
+    out = []
+    for start, end in segments_list:
+        out.append(a[:, start:end] @ b[start:end, :])
+    return mx.stack(out, axis=0)
+
+
+def bench_mlx(a, b, segments, warmup, iters):
+    for _ in range(warmup):
+        y = mx.segmented_mm(a, b, segments)
+        mx.eval(y)
+    mx.synchronize()
+
+    start = time.perf_counter()
+    for _ in range(iters):
+        y = mx.segmented_mm(a, b, segments)
+        mx.eval(y)
+    mx.synchronize()
+    end = time.perf_counter()
+    return (end - start) * 1e3 / iters
+
+
+def bench_mlx_loop(a, b, segments, warmup, iters):
+    for _ in range(warmup):
+        y = mlx_segmented_mm_loop(a, b, segments)
+        mx.eval(y)
+    mx.synchronize()
+
+    start = time.perf_counter()
+    for _ in range(iters):
+        y = mlx_segmented_mm_loop(a, b, segments)
+        mx.eval(y)
+    mx.synchronize()
+    end = time.perf_counter()
+    return (end - start) * 1e3 / iters
+
+
+def print_table(headers, rows):
+    widths = [len(h) for h in headers]
+    for row in rows:
+        for i, cell in enumerate(row):
+            widths[i] = max(widths[i], len(cell))
+
+    def fmt_row(row):
+        return (
+            "| "
+            + " | ".join(f"{cell:<{widths[i]}}" for i, cell in enumerate(row))
+            + " |"
+        )
+
+    sep = "|-" + "-|-".join("-" * w for w in widths) + "-|"
+    print(fmt_row(headers))
+    print(sep)
+    for row in rows:
+        print(fmt_row(row))
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--cases",
+        default=(
+            "128x128x1024x16,"
+            "128x128x1024x32,"
+            "256x256x2048x16,"
+            "512x512x4096x32,"
+            "1024x1024x4096x32,"
+            "1024x1024x8192x64"
+        ),
+        help="Comma-separated MxNxKxS list.",
+    )
+    parser.add_argument(
+        "--dtype",
+        default="float32",
+        choices=["float16", "bfloat16", "float32"],
+    )
+    parser.add_argument("--warmup", type=int, default=10)
+    parser.add_argument("--iters", type=int, default=50)
+    parser.add_argument(
+        "--segments",
+        choices=["equal", "random"],
+        default="random",
+        help="Segment generation pattern.",
+    )
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--no-check", action="store_true")
+    args = parser.parse_args()
+
+    mlx_dtype = MLX_DTYPES[args.dtype]
+
+    print(
+        f"dtype={args.dtype} warmup={args.warmup} iters={args.iters} segments={args.segments}"
+    )
+
+    headers = [
+        "Case",
+        "MLX ms",
+        "Loop ms",
+        "Speedup",
+        "MLX err",
+        "Loop err",
+    ]
+    rows = []
+
+    cases = parse_cases(args.cases)
+    for idx, (m, n, k, s) in enumerate(cases):
+        rng = np.random.default_rng(args.seed + idx)
+        a_np = rng.standard_normal((m, k)).astype(np.float32)
+        b_np = rng.standard_normal((k, n)).astype(np.float32)
+        seg_np = make_segments(k, s, args.segments, args.seed + idx)
+
+        a_mx = mx.array(a_np, dtype=mlx_dtype)
+        b_mx = mx.array(b_np, dtype=mlx_dtype)
+        seg_mx = mx.array(seg_np, dtype=mx.uint32)
+        mx.eval(a_mx, b_mx, seg_mx)
+
+        mlx_err_str = ""
+        loop_err_str = ""
+        if not args.no_check:
+            y_mlx = mx.segmented_mm(a_mx, b_mx, seg_mx)
+            y_loop = mlx_segmented_mm_loop(a_mx, b_mx, seg_mx)
+            mx.eval(y_mlx, y_loop)
+
+            if args.dtype == "float32":
+                ref = numpy_segmented_mm_ref(
+                    a_np.astype(np.float64),
+                    b_np.astype(np.float64),
+                    seg_np.tolist(),
+                )
+                mlx_err = np.max(np.abs(np.array(y_mlx, dtype=np.float64) - ref))
+                loop_err = np.max(np.abs(np.array(y_loop, dtype=np.float64) - ref))
+            else:
+                a_mx_f32 = mx.array(a_np, dtype=mx.float32)
+                b_mx_f32 = mx.array(b_np, dtype=mx.float32)
+                ref = mx.segmented_mm(a_mx_f32, b_mx_f32, seg_mx)
+                mx.eval(ref)
+                mlx_err = float(mx.max(mx.abs(ref - y_mlx.astype(mx.float32))).item())
+                loop_err = float(mx.max(mx.abs(ref - y_loop.astype(mx.float32))).item())
+            mlx_err_str = f"{mlx_err:.2e}"
+            loop_err_str = f"{loop_err:.2e}"
+
+        t_mlx = bench_mlx(a_mx, b_mx, seg_mx, args.warmup, args.iters)
+        t_loop = bench_mlx_loop(a_mx, b_mx, seg_mx, args.warmup, args.iters)
+        ratio = t_loop / t_mlx if t_mlx > 0 else float("inf")
+        rows.append(
+            [
+                f"{m}x{n}x{k}x{s}",
+                f"{t_mlx:.3f}",
+                f"{t_loop:.3f}",
+                f"{ratio:.2f}x",
+                mlx_err_str,
+                loop_err_str,
+            ]
+        )
+
+    print_table(headers, rows)
+    if not args.no_check:
+        if args.dtype == "float32":
+            print("err: max|result - numpy_fp64_ref|")
+        else:
+            print("err: max|result - own_fp32_result|")
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,177 @@
+# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Modified from
+# https://github.com/NVIDIA/cudnn-frontend/blob/main/cmake/cuDNN.cmake
+
+# Return the last file matching the pattern.
+function(find_file_glob VAR PATTERN)
+  file(GLOB _RESULT "${PATTERN}")
+  if(_RESULT)
+    list(LENGTH ${_RESULT} _RESULT_LENGTH)
+    if(_RESULT_LENGTH GREATER 0)
+      list(GET ${_RESULT} -1 _RESULT)
+    endif()
+    set(${VAR}
+        "${_RESULT}"
+        PARENT_SCOPE)
+  endif()
+endfunction()
+
+# Find the dir including the "cudnn.h" file.
+find_path(
+  CUDNN_INCLUDE_DIR cudnn.h
+  HINTS ${CUDNN_INCLUDE_PATH} ${CUDAToolkit_INCLUDE_DIRS}
+  PATH_SUFFIXES include OPTIONAL)
+
+# Glob searching "cudnn.h" for Windows.
+if(WIN32 AND NOT CUDNN_INCLUDE_DIR)
+  find_file_glob(
+    CUDNN_H_PATH
+    "C:/Program Files/NVIDIA/CUDNN/*/include/${CUDAToolkit_VERSION_MAJOR}.*/cudnn.h"
+  )
+  if(CUDNN_H_PATH)
+    get_filename_component(CUDNN_INCLUDE_DIR "${CUDNN_H_PATH}" DIRECTORY)
+  endif()
+endif()
+
+if(NOT CUDNN_INCLUDE_DIR)
+  message(
+    FATAL_ERROR
+      "Unable to find cudnn.h, please make sure cuDNN is installed and pass CUDNN_INCLUDE_PATH to cmake."
+  )
+endif()
+
+# Get cudnn version.
+file(READ "${CUDNN_INCLUDE_DIR}/cudnn_version.h" cudnn_version_header)
+string(REGEX MATCH "#define CUDNN_MAJOR [1-9]+" macrodef
+             "${cudnn_version_header}")
+string(REGEX MATCH "[1-9]+" CUDNN_MAJOR_VERSION "${macrodef}")
+
+# Function for searching library files.
+function(find_cudnn_library NAME)
+  if(NOT "${ARGV1}" STREQUAL "OPTIONAL")
+    set(_CUDNN_REQUIRED TRUE)
+  else()
+    set(_CUDNN_REQUIRED FALSE)
+  endif()
+
+  find_library(
+    ${NAME}_LIBRARY
+    NAMES ${NAME} "lib${NAME}.so.${CUDNN_MAJOR_VERSION}" NAMES_PER_DIR
+    HINTS ${CUDNN_LIBRARY_PATH} ${CUDAToolkit_LIBRARY_DIR}
+    PATH_SUFFIXES lib64 lib/x64 lib OPTIONAL)
+
+  if(WIN32 AND NOT ${NAME}_LIBRARY)
+    find_file_glob(
+      ${NAME}_LIBRARY
+      "C:/Program Files/NVIDIA/CUDNN/*/lib/${CUDAToolkit_VERSION_MAJOR}.*/x64/${NAME}.lib"
+    )
+  endif()
+
+  if(NOT ${NAME}_LIBRARY AND ${_CUDNN_REQUIRED})
+    message(
+      FATAL_ERROR
+        "Unable to find ${NAME}, please make sure cuDNN is installed and pass CUDNN_LIBRARY_PATH to cmake."
+    )
+  endif()
+
+  if(${NAME}_LIBRARY)
+    add_library(CUDNN::${NAME} UNKNOWN IMPORTED)
+    set_target_properties(
+      CUDNN::${NAME}
+      PROPERTIES INTERFACE_INCLUDE_DIRECTORIES ${CUDNN_INCLUDE_DIR}
+                 IMPORTED_LOCATION ${${NAME}_LIBRARY})
+    set(${NAME}_LIBRARY
+        "${${NAME}_LIBRARY}"
+        PARENT_SCOPE)
+  else()
+    message(STATUS "${NAME} not found.")
+  endif()
+endfunction()
+
+# Search for the main cudnn library.
+find_cudnn_library(cudnn)
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(CUDNN REQUIRED_VARS CUDNN_INCLUDE_DIR
+                                                      cudnn_LIBRARY)
+
+if(CUDNN_INCLUDE_DIR AND cudnn_LIBRARY)
+  set(CUDNN_FOUND
+      ON
+      CACHE INTERNAL "cuDNN Library Found")
+else()
+  set(CUDNN_FOUND
+      OFF
+      CACHE INTERNAL "cuDNN Library Not Found")
+endif()
+
+# Find out all the DLL files for Windows.
+if(WIN32 AND cudnn_LIBRARY)
+  get_filename_component(CUDNN_BIN_DIR "${cudnn_LIBRARY}" DIRECTORY)
+  string(REPLACE "/lib/" "/bin/" CUDNN_BIN_DIR "${CUDNN_BIN_DIR}")
+  file(
+    GLOB CUDNN_DLL_NAMES
+    RELATIVE "${CUDNN_BIN_DIR}"
+    "${CUDNN_BIN_DIR}/*.dll")
+endif()
+
+# Create an interface library that users can link with.
+add_library(CUDNN::cudnn_all INTERFACE IMPORTED)
+target_link_libraries(CUDNN::cudnn_all INTERFACE CUDNN::cudnn)
+target_include_directories(
+  CUDNN::cudnn_all INTERFACE $<INSTALL_INTERFACE:include>
+                             $<BUILD_INTERFACE:${CUDNN_INCLUDE_DIR}>)
+
+# Add other components of cudnn.
+if(CUDNN_MAJOR_VERSION EQUAL 8)
+  find_cudnn_library(cudnn_adv_infer)
+  find_cudnn_library(cudnn_adv_train)
+  find_cudnn_library(cudnn_cnn_infer)
+  find_cudnn_library(cudnn_cnn_train)
+  find_cudnn_library(cudnn_ops_infer)
+  find_cudnn_library(cudnn_ops_train)
+
+  target_link_libraries(
+    CUDNN::cudnn_all
+    INTERFACE CUDNN::cudnn_adv_train CUDNN::cudnn_ops_train
+              CUDNN::cudnn_cnn_train CUDNN::cudnn_adv_infer
+              CUDNN::cudnn_cnn_infer CUDNN::cudnn_ops_infer)
+
+elseif(CUDNN_MAJOR_VERSION EQUAL 9)
+  find_cudnn_library(cudnn_graph)
+  find_cudnn_library(cudnn_engines_runtime_compiled)
+  find_cudnn_library(cudnn_ops OPTIONAL)
+  find_cudnn_library(cudnn_cnn OPTIONAL)
+  find_cudnn_library(cudnn_adv OPTIONAL)
+  find_cudnn_library(cudnn_engines_precompiled OPTIONAL)
+  find_cudnn_library(cudnn_heuristic OPTIONAL)
+
+  target_link_libraries(
+    CUDNN::cudnn_all
+    INTERFACE CUDNN::cudnn_graph
+              CUDNN::cudnn_engines_runtime_compiled
+              CUDNN::cudnn_ops
+              CUDNN::cudnn_cnn
+              CUDNN::cudnn_adv
+              CUDNN::cudnn_engines_precompiled
+              CUDNN::cudnn_heuristic)
+endif()
@@ -26,6 +26,7 @@ ENABLE_PREPROCESSING   = YES
 MACRO_EXPANSION        = YES
 EXPAND_ONLY_PREDEF     = NO
 SKIP_FUNCTION_MACROS   = NO
+PREDEFINED             = MLX_API=

 ################################################################################
 # Compound extraction control.                                                 #
@@ -38,3 +38,17 @@ the docs. Then force add the `build/html` directory:
 `git add -f build/html`

 Commit and push the changes to the `gh-pages` branch.
+
+## Doc Development Setup
+
+To enable live refresh of docs while writing:
+
+Install sphinx autobuild
+```
+pip install sphinx-autobuild
+```
+
+Run auto build on docs/src folder
+```
+sphinx-autobuild ./src ./build/html
+```
@@ -45,7 +45,7 @@ The next step is to setup a CMake file in ``CMakeLists.txt``:

  project(example LANGUAGES CXX)

-  set(CMAKE_CXX_STANDARD 17)
+  set(CMAKE_CXX_STANDARD 20)
  set(CMAKE_CXX_STANDARD_REQUIRED ON)


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

 .. toctree::
   :caption: Python API Reference
@@ -83,7 +83,8 @@ Build from source
 Build Requirements
 ^^^^^^^^^^^^^^^^^^

- A C++ compiler with C++17 support (e.g. Clang >= 5.0)
+- ``libblas-dev``, ``liblapack-dev``, and ``liblapacke-dev`` (Linux)
+- A C++ compiler with C++20 support (e.g. Clang >= 15.0)
 - `cmake <https://cmake.org/>`_ -- version 3.25 or later, and ``make``
 - Xcode >= 15.0 and macOS SDK >= 14.0

@@ -17,3 +17,5 @@ Devices and Streams
   set_default_stream
   stream
   synchronize
+   device_count
+   device_info
@@ -175,6 +175,7 @@ In detail:
   value_and_grad
   quantize
   average_gradients
+   fsdp_apply_gradients

 .. toctree::

@@ -183,3 +184,4 @@ In detail:
   nn/functions
   nn/losses
   nn/init
+   nn/distributed
@@ -0,0 +1,30 @@
+.. _nn_distributed:
+
+Distributed
+-----------
+
+Helper Routines
+^^^^^^^^^^^^^^^
+
+The :code:`mlx.nn.layers.distributed` package contains helpful routines to 
+create sharded layers from existing :class:`Modules <mlx.nn.Module>`.
+
+.. currentmodule:: mlx.nn.layers.distributed
+.. autosummary::
+   :toctree: _autosummary
+
+   shard_linear
+   shard_inplace
+
+Layers
+^^^^^^
+
+.. currentmodule:: mlx.nn
+.. autosummary::
+   :toctree: _autosummary
+   :template: nn-module-template.rst
+
+   AllToShardedLinear
+   ShardedToAllLinear
+   QuantizedAllToShardedLinear
+   QuantizedShardedToAllLinear
@@ -10,6 +10,7 @@ Layers
   :template: nn-module-template.rst

   ALiBi
+   AllToShardedLinear
   AvgPool1d
   AvgPool2d
   AvgPool3d
@@ -46,8 +47,10 @@ Layers
   Mish
   MultiHeadAttention
   PReLU
+   QuantizedAllToShardedLinear
   QuantizedEmbedding
   QuantizedLinear
+   QuantizedShardedToAllLinear
   RMSNorm
   ReLU
   ReLU2
@@ -56,6 +59,7 @@ Layers
   RoPE
   SELU
   Sequential
+   ShardedToAllLinear
   Sigmoid
   SiLU
   SinusoidalPositionalEncoding
@@ -11,6 +11,7 @@ Transforms
   eval
   async_eval
   compile
+   checkpoint
   custom_function
   disable_compile
   enable_compile
@@ -117,89 +117,11 @@ The following examples aim to clarify the backend initialization logic in MLX:
    world_ring = mx.distributed.init(backend="ring")
    world_any = mx.distributed.init()  # same as MPI because it was initialized first!

-.. _training_example:
+Distributed Program Examples
+----------------------------

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

 .. _ring_section:

@@ -155,6 +155,34 @@ parameters, pass them as inputs to the ``call`` wrapper:
   mx.export_function("model.mlxfn", call, (mx.zeros(4),), params)


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

@@ -22,16 +22,57 @@ target_sources(
 # Define MLX_VERSION only in the version.cpp file.
 add_library(mlx_version OBJECT ${CMAKE_CURRENT_SOURCE_DIR}/version.cpp)
 target_compile_definitions(mlx_version PRIVATE MLX_VERSION="${MLX_VERSION}")
+target_include_directories(mlx_version PRIVATE ${PROJECT_SOURCE_DIR})
 target_link_libraries(mlx PRIVATE $<BUILD_INTERFACE:mlx_version>)

-if(MSVC)
-  # Disable some MSVC warnings to speed up compilation.
-  target_compile_options(mlx PUBLIC /wd4068 /wd4244 /wd4267 /wd4804)
+# Do not export symbols by default.
+set_target_properties(
+  mlx mlx_version
+  PROPERTIES VISIBILITY_INLINES_HIDDEN ON
+             CXX_VISIBILITY_PRESET hidden
+             CUDA_VISIBILITY_PRESET hidden)
+
+# Define MLX_EXPORT for shared libraries, MLX_STATIC for static libraries.
+set_target_properties(mlx PROPERTIES DEFINE_SYMBOL MLX_EXPORT)
+if(BUILD_SHARED_LIBS)
+  target_compile_definitions(mlx_version PUBLIC MLX_EXPORT)
+else()
+  target_compile_definitions(mlx PUBLIC MLX_STATIC)
+  target_compile_definitions(mlx_version PUBLIC MLX_STATIC)
 endif()

-if(WIN32)
-  # Export symbols by default to behave like macOS/linux.
-  set_target_properties(mlx PROPERTIES WINDOWS_EXPORT_ALL_SYMBOLS TRUE)
+if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
+  # Supress warnings: note: parameter passing for argument of type
+  # 'std::pair<float, float>' when C++17 is enabled changed to match C++14 in
+  # GCC 10.1
+  target_compile_options(mlx PRIVATE -Wno-psabi)
+endif()
+
+if(MSVC)
+  # Some of CUDA's headers include windows.h, which defines min/max macros.
+  target_compile_definitions(mlx PRIVATE NOMINMAX WIN32_LEAN_AND_MEAN)
+  # Unicode support in fmt does not compile in .cu files.
+  target_compile_definitions(mlx PRIVATE FMT_UNICODE=0)
+  # Disable some MSVC warnings to speed up compilation.
+  target_compile_options(
+    mlx
+    PUBLIC $<$<COMPILE_LANGUAGE:CXX>:/wd4244 /wd4267>
+    PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/wd4068
+            /wd4146
+            /wd4700
+            /wd4804
+            /wd4805>
+            $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/wd4244
+            -Xcompiler=/wd4267>)
+  # Enable /bigobj for heavily templated code (e.g., binary.cpp) that exceeds
+  # the default 65,535 section limit in COFF object files.
+  target_compile_options(
+    mlx PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/bigobj>
+                $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/bigobj>)
+  # Use modern preprocessor, otherwise CCCL would complain.
+  target_compile_options(
+    mlx PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/Zc:preprocessor>
+                $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/Zc:preprocessor>)
 endif()

 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/common)
@@ -4,12 +4,14 @@

 #include <cstdlib>

+#include "mlx/api.h"
+
 namespace mlx::core::allocator {

 // Simple wrapper around buffer pointers
 // WARNING: Only Buffer objects constructed from and those that wrap
 //          raw pointers from mlx::allocator are supported.
-class Buffer {
+class MLX_API Buffer {
 private:
  void* ptr_;

@@ -28,7 +30,7 @@ class Buffer {
  };
 };

-class Allocator {
+class MLX_API Allocator {
  /** Abstract base class for a memory allocator. */
 public:
  virtual Buffer malloc(size_t size) = 0;
@@ -47,7 +49,7 @@ class Allocator {
  virtual ~Allocator() = default;
 };

-Allocator& allocator();
+MLX_API Allocator& allocator();

 inline Buffer malloc(size_t size) {
  return allocator().malloc(size);
@@ -0,0 +1,29 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+// MLX_API macro for controlling symbol visibility, must add for public APIs.
+//
+// Usage:
+//   MLX_API void some_function(...);
+//   class MLX_API SomeClass { ... };
+
+#if defined(MLX_STATIC)
+
+// Static library build - no import/export decorations needed
+#define MLX_API
+
+#else
+
+// Shared library build.
+#if defined(_WIN32)
+#if defined(MLX_EXPORT)
+#define MLX_API __declspec(dllexport)
+#else
+#define MLX_API __declspec(dllimport)
+#endif // defined(MLX_EXPORT)
+#else
+#define MLX_API __attribute__((visibility("default")))
+#endif // defined(_WIN32)
+
+#endif // defined(MLX_STATIC)
@@ -21,11 +21,12 @@ array::array(
    Dtype dtype,
    std::shared_ptr<Primitive> primitive,
    std::vector<array> inputs)
-    : array_desc_(std::make_shared<ArrayDesc>(
-          std::move(shape),
-          dtype,
-          std::move(primitive),
-          std::move(inputs))) {
+    : array_desc_(
+          std::make_shared<ArrayDesc>(
+              std::move(shape),
+              dtype,
+              std::move(primitive),
+              std::move(inputs))) {
  if (has_primitive() && this->primitive().stream().device == Device::gpu) {
    for (auto& in : this->inputs()) {
      if (in.dtype() == float64) {
@@ -69,16 +70,18 @@ array array::unsafe_weak_copy(const array& other) {
 }

 array::array(std::initializer_list<float> data)
-    : array_desc_(std::make_shared<ArrayDesc>(
-          Shape{static_cast<ShapeElem>(data.size())},
-          float32)) {
+    : array_desc_(
+          std::make_shared<ArrayDesc>(
+              Shape{static_cast<ShapeElem>(data.size())},
+              float32)) {
  init(data.begin());
 }

 array::array(std::initializer_list<int> data, Dtype dtype)
-    : array_desc_(std::make_shared<ArrayDesc>(
-          Shape{static_cast<ShapeElem>(data.size())},
-          dtype)) {
+    : array_desc_(
+          std::make_shared<ArrayDesc>(
+              Shape{static_cast<ShapeElem>(data.size())},
+              dtype)) {
  init(data.begin());
 }

@@ -131,6 +134,7 @@ bool array::is_available() const {
  } else if (
      status() == Status::evaluated &&
      (!event().valid() || event().is_signaled())) {
+    detach_event();
    set_status(Status::available);
    return true;
  }
@@ -8,6 +8,7 @@
 #include <vector>

 #include "mlx/allocator.h"
+#include "mlx/api.h"
 #include "mlx/dtype.h"
 #include "mlx/event.h"
 #include "mlx/small_vector.h"
@@ -22,7 +23,7 @@ using ShapeElem = int32_t;
 using Shape = SmallVector<ShapeElem>;
 using Strides = SmallVector<int64_t>;

-class array {
+class MLX_API array {
  /* An array is really a node in a graph. It contains a shared ArrayDesc
   * object */

@@ -121,7 +122,7 @@ class array {
   *  This function supports negative indexing and provides
   *  bounds checking. */
  auto shape(int dim) const {
-    return shape().at(dim < 0 ? dim + ndim() : dim);
+    return shape().at(dim < 0 ? dim + static_cast<int>(ndim()) : dim);
  }

  /** The strides of the array. */
@@ -135,7 +136,7 @@ class array {
   *  This function supports negative indexing and provides
   *  bounds checking. */
  auto strides(int dim) const {
-    return strides().at(dim < 0 ? dim + ndim() : dim);
+    return strides().at(dim < 0 ? dim + static_cast<int>(ndim()) : dim);
  }

  /** Get the arrays data type. */
@@ -153,7 +154,7 @@ class array {
  template <typename T>
  T item() const;

-  struct ArrayIterator {
+  struct MLX_API ArrayIterator {
    using iterator_category = std::random_access_iterator_tag;
    using difference_type = size_t;
    using value_type = const array;
@@ -464,7 +465,7 @@ class array {
  template <typename It>
  void init(const It src);

-  struct ArrayDesc {
+  struct MLX_API ArrayDesc {
    Shape shape;
    Strides strides;
    size_t size;
@@ -488,10 +489,10 @@ class array {
    int64_t offset{0};

    // The size in elements of the data buffer the array accesses
-    size_t data_size;
+    size_t data_size{0};

    // Contains useful meta data about the array
-    Flags flags;
+    Flags flags{true, true, true};

    std::vector<array> inputs;
    // An array to keep track of the siblings from a multi-output
@@ -541,9 +542,10 @@ template <typename T>
 array::array(
    std::initializer_list<T> data,
    Dtype dtype /* = TypeToDtype<T>() */)
-    : array_desc_(std::make_shared<ArrayDesc>(
-          Shape{static_cast<ShapeElem>(data.size())},
-          dtype)) {
+    : array_desc_(
+          std::make_shared<ArrayDesc>(
+              Shape{static_cast<ShapeElem>(data.size())},
+              dtype)) {
  init(data.begin());
 }

@@ -2,6 +2,7 @@

 #pragma once

+#include <algorithm>
 #include <cassert>
 #include <functional>
 #include <map>
@@ -0,0 +1,14 @@
+// Copyright © 2026 Apple Inc.
+
+namespace mlx::core {
+
+inline constexpr short get_pack_factor(int bits, int wsize = 8) {
+  return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits);
+}
+
+inline constexpr short get_bytes_per_pack(int bits, int wsize = 8) {
+  bool power_of_2_bits = (bits & (bits - 1)) == 0;
+  return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3);
+}
+
+} // namespace mlx::core
@@ -40,7 +40,7 @@ add_dependencies(mlx cpu_compiled_preamble)

 target_sources(
  mlx
-  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/available.cpp
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/arg_reduce.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/binary.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp
@@ -1,11 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/cpu/available.h"
-
-namespace mlx::core::cpu {
-
-bool is_available() {
-  return true;
-}
-
-} // namespace mlx::core::cpu
@@ -1,9 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-namespace mlx::core::cpu {
-
-bool is_available();
-
-} // namespace mlx::core::cpu
@@ -119,13 +119,15 @@ void* compile(
    source_file.close();

    try {
-      JitCompiler::exec(JitCompiler::build_command(
-          output_dir, source_file_name, shared_lib_name));
+      JitCompiler::exec(
+          JitCompiler::build_command(
+              output_dir, source_file_name, shared_lib_name));
    } catch (const std::exception& error) {
-      throw std::runtime_error(fmt::format(
-          "[Compile::eval_cpu] Failed to compile function {0}: {1}",
-          kernel_name,
-          error.what()));
+      throw std::runtime_error(
+          fmt::format(
+              "[Compile::eval_cpu] Failed to compile function {0}: {1}",
+              kernel_name,
+              error.what()));
    }
  }

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

-struct CommandEncoder {
+struct MLX_API CommandEncoder {
  CommandEncoder(Stream stream) : stream_(stream) {}

  CommandEncoder(const CommandEncoder&) = delete;
@@ -62,6 +62,6 @@ struct CommandEncoder {
  int num_ops_{0};
 };

-CommandEncoder& get_command_encoder(Stream stream);
+MLX_API CommandEncoder& get_command_encoder(Stream stream);

 } // namespace mlx::core::cpu
@@ -761,7 +761,7 @@ void masked_scatter_impl(const array& mask, const array& src, array& out) {
  const size_t mask_batch_size = mask.size() / batch_count;
  const size_t src_batch_size = src.size() / batch_count;

-  for (uint b = 0; b < batch_count; ++b) {
+  for (size_t b = 0; b < batch_count; ++b) {
    size_t src_consumed = 0;
    src_it.seek(b * src_batch_size);

@@ -34,18 +34,30 @@ struct VisualStudioInfo {
    arch = "x64";
 #endif
    // Get path of Visual Studio.
-    std::string vs_path = JitCompiler::exec(fmt::format(
-        "\"{0}\\Microsoft Visual Studio\\Installer\\vswhere.exe\""
-        " -property installationPath",
-        std::getenv("ProgramFiles(x86)")));
+    // Use -latest to get only the most recent installation when multiple
+    // versions are installed, avoiding path concatenation issues.
+    std::string vs_path = JitCompiler::exec(
+        fmt::format(
+            "\"{0}\\Microsoft Visual Studio\\Installer\\vswhere.exe\""
+            " -latest -property installationPath",
+            std::getenv("ProgramFiles(x86)")));
    if (vs_path.empty()) {
      throw std::runtime_error("Can not find Visual Studio.");
    }
+    // Trim any trailing whitespace/newlines from the path
+    vs_path.erase(
+        std::find_if(
+            vs_path.rbegin(),
+            vs_path.rend(),
+            [](unsigned char ch) { return !std::isspace(ch); })
+            .base(),
+        vs_path.end());
    // Read the envs from vcvarsall.
-    std::string envs = JitCompiler::exec(fmt::format(
-        "\"{0}\\VC\\Auxiliary\\Build\\vcvarsall.bat\" {1} >NUL && set",
-        vs_path,
-        arch));
+    std::string envs = JitCompiler::exec(
+        fmt::format(
+            "\"{0}\\VC\\Auxiliary\\Build\\vcvarsall.bat\" {1} >NUL && set",
+            vs_path,
+            arch));
    for (const std::string& line : str_split(envs, '\n')) {
      // Each line is in the format "ENV_NAME=values".
      auto pos = line.find_first_of('=');
@@ -140,12 +152,13 @@ std::string JitCompiler::exec(const std::string& cmd) {
  int code = WEXITSTATUS(status);
 #endif
  if (code != 0) {
-    throw std::runtime_error(fmt::format(
-        "Failed to execute command with return code {0}: \"{1}\", "
-        "the output is: {2}",
-        code,
-        cmd,
-        ret));
+    throw std::runtime_error(
+        fmt::format(
+            "Failed to execute command with return code {0}: \"{1}\", "
+            "the output is: {2}",
+            code,
+            cmd,
+            ret));
  }
  return ret;
 }
@@ -1,5 +1,6 @@
 // Copyright © 2023 Apple Inc.

+#include "mlx/backend/common/quantized.h"
 #include "mlx/backend/common/unary.h"
 #include "mlx/backend/cpu/copy.h"
 #include "mlx/backend/cpu/encoder.h"
@@ -14,6 +15,19 @@ namespace mlx::core {

 namespace {

+array ensure_row_contiguous(
+    const array& arr,
+    cpu::CommandEncoder& encoder,
+    Stream s) {
+  if (arr.flags().row_contiguous) {
+    return arr;
+  } else {
+    auto arr_cpy = contiguous_copy_cpu(arr, s);
+    encoder.add_temporary(arr_cpy);
+    return arr_cpy;
+  }
+};
+
 const static float FP4_LUT[16] = {
    +0.0f,
    +0.5f,
@@ -47,15 +61,6 @@ static inline T dequantize_scale(uint8_t s) {
  }
 }

-inline constexpr short get_pack_factor(int bits, int wsize = 8) {
-  return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits);
-}
-
-inline constexpr short get_bytes_per_pack(int bits, int wsize = 8) {
-  auto power_of_2_bits = (bits & (bits - 1)) == 0;
-  return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3);
-}
-
 template <typename T, int bits>
 void extract_bits(const uint8_t* w_in, T* w_out) {
  static_assert(bits == 3 || bits == 5 || bits == 6);
@@ -922,20 +927,9 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
  auto& scales_pre = inputs[2];

  auto& encoder = cpu::get_command_encoder(stream());
-  auto ensure_row_contiguous = [s = stream(), &encoder](const array& arr) {
-    if (arr.flags().row_contiguous) {
-      return arr;
-    } else {
-      auto arr_cpy = array(arr.shape(), arr.dtype(), nullptr, {});
-      copy_cpu(arr, arr_cpy, CopyType::General, s);
-      encoder.add_temporary(arr_cpy);
-      return arr_cpy;
-    }
-  };
-
-  auto x = ensure_row_contiguous(x_pre);
-  auto w = ensure_row_contiguous(w_pre);
-  auto scales = ensure_row_contiguous(scales_pre);
+  auto x = ensure_row_contiguous(x_pre, encoder, stream());
+  auto w = ensure_row_contiguous(w_pre, encoder, stream());
+  auto scales = ensure_row_contiguous(scales_pre, encoder, stream());

  out.set_data(allocator::malloc(out.nbytes()));

@@ -944,7 +938,7 @@ void QuantizedMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
  encoder.set_input_array(scales);
  encoder.set_output_array(out);
  if (mode_ == QuantizationMode::Affine) {
-    auto biases = ensure_row_contiguous(inputs[3]);
+    auto biases = ensure_row_contiguous(inputs[3], encoder, stream());
    encoder.set_input_array(biases);
    encoder.dispatch([out = array::unsafe_weak_copy(out),
                      x = array::unsafe_weak_copy(x),
@@ -1052,6 +1046,105 @@ void GatherQMM::eval_cpu(const std::vector<array>& inputs, array& out) {
  }
 }

+uint8_t to_fp8_e8m0(float x) {
+  if (!std::isfinite(x)) {
+    return 0xFF;
+  }
+  if (x < 0.0f) {
+    return 0x00;
+  }
+  float le = std::log2(x);
+  int n = int(std::round(le));
+
+  n = n < -127 ? -127 : n;
+  n = n > 127 ? 127 : n;
+  return static_cast<uint8_t>(n + 127);
+}
+
+uint8_t to_fp4_e2m1(float x) {
+  if (std::isnan(x)) {
+    return 0x7;
+  }
+
+  const uint8_t sign_bit = (std::signbit(x)) ? 0x8 : 0x0;
+  x = std::abs(x);
+
+  uint8_t bits;
+  if (x > 5.0f) {
+    bits = 0x7;
+  } else if (x >= 3.5f) {
+    bits = 0x6;
+  } else if (x > 2.5f) {
+    bits = 0x5;
+  } else if (x >= 1.75f) {
+    bits = 0x4;
+  } else if (x > 1.25f) {
+    bits = 0x3;
+  } else if (x >= 0.75f) {
+    bits = 0x2;
+  } else if (x > 0.25f) {
+    bits = 0x1;
+  } else {
+    bits = 0x0;
+  }
+  return bits | sign_bit;
+}
+
+template <typename T>
+void fp_quantize_dequantize(
+    const array& w_arr,
+    array& out_arr,
+    int bits,
+    int group_size,
+    size_t w_size) {
+  auto w = w_arr.data<T>();
+  auto out = out_arr.data<T>();
+
+  size_t n_groups = w_size / group_size;
+
+  for (size_t i = 0; i < n_groups; ++i) {
+    size_t idx = i * group_size;
+    float scale = -std::numeric_limits<float>::infinity();
+    for (int j = 0; j < group_size; ++j) {
+      scale = std::max(scale, std::abs(w[idx + j]));
+    }
+    scale /= bits == 4 ? 6.0f : 448.0f;
+    if (group_size == 16) {
+      scale = dequantize_scale<float, 16>(detail::ToFP8()(scale));
+    } else {
+      scale = dequantize_scale<float, 32>(to_fp8_e8m0(scale));
+    }
+
+    for (int j = 0; j < group_size; ++j) {
+      float w_el = scale == 0 ? 0.0f : w[idx + j] / scale;
+      float output;
+      if (bits == 8) {
+        output = detail::FromFP8()(detail::ToFP8()(w_el));
+      } else {
+        output = FP4_LUT[to_fp4_e2m1(w_el)];
+      }
+      out[idx + j] = static_cast<T>(scale * output);
+    }
+  }
+}
+
+void dispatch_quantize_dequantize(
+    const array& w,
+    array& out,
+    int bits,
+    int group_size) {
+  if (w.dtype() == float16) {
+    fp_quantize_dequantize<float16_t>(w, out, bits, group_size, w.size());
+  } else if (w.dtype() == bfloat16) {
+    fp_quantize_dequantize<bfloat16_t>(w, out, bits, group_size, w.size());
+  } else if (w.dtype() == float32) {
+    fp_quantize_dequantize<float>(w, out, bits, group_size, w.size());
+  } else {
+    throw std::runtime_error(
+        "[quantize_dequantize] Only supports floating point inputs");
+  }
+}
+
 template <typename T, typename U>
 void quantize(
    const T* w,
@@ -1136,15 +1229,8 @@ void dispatch_quantize(
 void fast::Quantize::eval_cpu(
    const std::vector<array>& inputs,
    std::vector<array>& outputs) {
-  auto ensure_row_contiguous = [s = stream()](const array& arr) {
-    if (arr.flags().row_contiguous) {
-      return std::make_pair(arr, false);
-    } else {
-      return std::make_pair(contiguous_copy_cpu(arr, s), true);
-    }
-  };
-
-  auto [w, copied] = ensure_row_contiguous(inputs[0]);
+  auto& encoder = cpu::get_command_encoder(stream());
+  auto w = ensure_row_contiguous(inputs[0], encoder, stream());
  auto& out = outputs[0];
  out.set_data(allocator::malloc(out.nbytes()));

@@ -1152,10 +1238,6 @@ void fast::Quantize::eval_cpu(
  auto& biases = outputs[2];
  scales.set_data(allocator::malloc(scales.nbytes()));
  biases.set_data(allocator::malloc(biases.nbytes()));
-  auto& encoder = cpu::get_command_encoder(stream());
-  if (copied) {
-    encoder.add_temporary(w);
-  }
  encoder.set_input_array(w);
  encoder.set_input_array(scales);
  encoder.set_input_array(biases);
@@ -1238,6 +1320,43 @@ void fast::ConvertFP8::eval_cpu(
 }

 void QQMatmul::eval_cpu(const std::vector<array>& inputs, array& out) {
-  throw std::runtime_error("QQMatmul not implemented on CPU.");
+  auto& encoder = cpu::get_command_encoder(stream());
+
+  bool w_quantized = (inputs[1].dtype() == uint32);
+  if (w_quantized && inputs[0].shape(-2) == 1) {
+    bool donate_x = inputs[0].is_donatable();
+    auto x = ensure_row_contiguous(inputs[0], encoder, stream());
+    auto w = ensure_row_contiguous(inputs[1], encoder, stream());
+    auto scales = ensure_row_contiguous(inputs[2], encoder, stream());
+
+    out.set_data(allocator::malloc(out.nbytes()));
+
+    // If x is a copy it should be donatable
+    donate_x |= x.is_donatable();
+    auto xhat = donate_x
+        ? x
+        : array(allocator::malloc(x.nbytes()), x.shape(), x.dtype());
+    if (!donate_x) {
+      encoder.add_temporary(xhat);
+    }
+    encoder.set_input_array(x);
+    encoder.set_input_array(w);
+    encoder.set_input_array(scales);
+    encoder.set_output_array(out);
+    encoder.dispatch([out = array::unsafe_weak_copy(out),
+                      x = array::unsafe_weak_copy(x),
+                      xhat = array::unsafe_weak_copy(xhat),
+                      w = array::unsafe_weak_copy(w),
+                      scales = array::unsafe_weak_copy(scales),
+                      group_size_ = group_size_,
+                      bits_ = bits_]() mutable {
+      dispatch_quantize_dequantize(x, xhat, bits_, group_size_);
+      fp_qmm_dispatch(out, xhat, w, scales, group_size_, bits_, true);
+    });
+    return;
+  } else {
+    throw std::runtime_error("[QQMatmul] NYI for the general case");
+  }
 }
+
 } // namespace mlx::core
@@ -1,11 +1,18 @@
 #pragma once

+// Required for using M_LN2 in MSVC.
+#define _USE_MATH_DEFINES
+
+#include <math.h>
 #include <stdint.h>
 #include <algorithm>
-#include <cmath>
 #include <complex>
 #include <functional>

+#ifdef _MSC_VER
+#include <intrin.h> // For _BitScanReverse
+#endif
+
 namespace mlx::core::simd {
 template <typename T, int N>
 struct Simd;
@@ -22,6 +29,14 @@ struct Simd<T, 1> {
  Simd(Simd<U, 1> v) : value(v.value) {}
  template <typename U>
  Simd(U v) : value(v) {}
+
+  T operator[](int) const {
+    return value;
+  }
+
+  T& operator[](int) {
+    return value;
+  }
 };

 template <typename T, int N>
@@ -105,7 +120,7 @@ Simd<T, 1> log1p(Simd<T, 1> in) {
      if (r == 0) { // handle underflow
        return Simd<T, 1>{T{x, theta}};
      }
-      return Simd<T, 1>{T{((typeof(x))(0.5)) * std::log1p(r), theta}};
+      return Simd<T, 1>{T{((decltype(x))(0.5)) * std::log1p(r), theta}};
    } else {
      auto z0 = std::hypot(x + 1, y);
      return Simd<T, 1>{T{std::log(z0), theta}};
@@ -173,7 +188,16 @@ DEFAULT_BINARY(||)

 template <typename T>
 Simd<T, 1> clz(Simd<T, 1> x_) {
+#ifdef _MSC_VER
+  // MSVC doesn't have __builtin_clz, use _BitScanReverse instead
+  unsigned long index;
+  if (_BitScanReverse(&index, static_cast<unsigned long>(x_.value))) {
+    return static_cast<T>(31 - index);
+  }
+  return static_cast<T>(32); // All zeros case
+#else
  return __builtin_clz(x_.value);
+#endif
 }

 template <typename T>
@@ -154,24 +154,19 @@ struct ToFP8 {
 struct FromFP8 {
  template <int N>
  Simd<float, N> operator()(Simd<uint8_t, N> x) {
-    auto w = Simd<uint32_t, N>(x) << 24;
-    auto sign = w & 0x80000000;
-    auto nonsign = w & 0x7FFFFFFF;
-
-    auto renorm_shift = clz(nonsign);
-    renorm_shift = simd::select(
-        renorm_shift > Simd<uint32_t, N>{4},
-        renorm_shift - Simd<uint32_t, N>{4},
-        Simd<uint32_t, N>{0});
-
-    Simd<int32_t, N> inf_nan_mask =
-        (Simd<int32_t, N>(nonsign + 0x01000000) >> 8) & 0x7F800000;
-    auto zero_mask = Simd<int32_t, N>(nonsign - 1) >> 31;
-    auto result = sign |
-        ((((nonsign << renorm_shift >> 4) + ((0x78 - renorm_shift) << 23)) |
-          inf_nan_mask) &
-         ~zero_mask);
-    return fp32_from_bits(result);
+    auto v = Simd<uint16_t, N>(x & 127) << 7;
+    Simd<float, N> out;
+    if constexpr (simd::max_size<float16_t> >= N) {
+      auto converted = *(Simd<float16_t, N>*)(&v);
+      out = converted * 256.0;
+    } else {
+      for (int i = 0; i < N; ++i) {
+        auto converted = *(float16_t*)(&v[i]);
+        out[i] = converted * 256.0;
+      }
+    }
+    auto sign = Simd<bool, N>(x & 128);
+    return select(sign, -out, out);
  }
  float operator()(uint8_t x) {
    return (*this)(Simd<uint8_t, 1>(x)).value;
@@ -19,8 +19,8 @@ target_sources(
          ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_conv.cu
          ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_grouped_conv.cu
          ${CMAKE_CURRENT_SOURCE_DIR}/cublas_utils.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/cuda.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/cudnn_utils.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/custom_kernel.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cu
@@ -30,6 +30,7 @@ target_sources(
          ${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}/hadamard.cu
          ${CMAKE_CURRENT_SOURCE_DIR}/jit_module.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/kernel_utils.cu
@@ -57,21 +58,24 @@ target_sources(
          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/affine_quantize.cu
          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/fp_quantize.cu
          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/quantized.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_utils.cu
          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/convert_fp8.cu
          ${CMAKE_CURRENT_SOURCE_DIR}/worker.cpp)

 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/binary)
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/quantized/qmm)
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/unary)

 # fp4 is not available on < 12.8
 if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 12.8.0)
  target_include_directories(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/)
+  target_sources(mlx
+                 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/no_qqmm_impl.cpp)
 else()
  target_sources(
-    mlx
-    PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm.cpp
-            ${CMAKE_CURRENT_SOURCE_DIR}/quantized/cublas_qqmm.cpp
-            ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_utils.cu)
+    mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_impl.cpp
+                ${CMAKE_CURRENT_SOURCE_DIR}/quantized/cublas_qqmm.cpp)
 endif()

 if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.9.0)
@@ -113,22 +117,22 @@ target_compile_options(mlx
 target_compile_options(
  mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>")

-# CUDA 12.8 emits warning #20280-D for copy kernels which is a false positive.
-# Explicitly pass this flag to suppress the warning, it is safe to set it to
-# true but the warning wouldn't be suppressed.
-if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8.0)
+# Required for generating optimized CUTLASS code.
+if(NOT MSVC)
  target_compile_options(
-    mlx
-    PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--static-global-template-stub=false>")
+    mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-fno-strict-aliasing>")
 endif()

-# Suppress warning when building for compute capability 7 used by V100.
+# Suppress nvcc warnings on C++ headers.
 target_compile_options(
-  mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--Wno-deprecated-gpu-targets>")
+  mlx
+  PRIVATE
+    $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe="--diag_suppress=27,997,1394,20011,20208">
+)

-# Suppress nvcc warnings on MLX headers.
-target_compile_options(mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe
-                                   --diag_suppress=997>)
+# Ignore some valid nvcc warnings, we might want to fix them in future.
+target_compile_options(
+  mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe="--diag_suppress=177,550">)

 # Use stronger binaries compression. This feature was introduced in CUDA 12.8
 # and requires drivers released after CUDA 12.4.
@@ -143,12 +147,11 @@ if(NOT DEFINED MLX_CUDA_ARCHITECTURES)
    COMMAND __nvcc_device_query
    OUTPUT_VARIABLE MLX_CUDA_ARCHITECTURES
    OUTPUT_STRIP_TRAILING_WHITESPACE)
-  set(UPGRADABLE_ARCHITECTURES "90;100;121")
  if(MLX_CUDA_ARCHITECTURES STREQUAL "")
    message(
      FATAL_ERROR
        "Can not get native CUDA arch, must set MLX_CUDA_ARCHITECTURES")
-  elseif(MLX_CUDA_ARCHITECTURES IN_LIST UPGRADABLE_ARCHITECTURES)
+  elseif(MLX_CUDA_ARCHITECTURES GREATER_EQUAL 90)
    # Use arch-specific compute capability whenever possible.
    set(MLX_CUDA_ARCHITECTURES "${MLX_CUDA_ARCHITECTURES}a")
  endif()
@@ -157,16 +160,53 @@ 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")
+# Skip Hopper-only kernels when not building for sm90a.
+if(NOT DEFINED ENV{MLX_DISABLE_SM90A_KERNELS}
+   AND (("90a" IN_LIST MLX_CUDA_ARCHITECTURES) OR ("90a-real" IN_LIST
+                                                   MLX_CUDA_ARCHITECTURES)))
+  target_compile_definitions(mlx PRIVATE MLX_CUDA_SM90A_ENABLED)
+endif()
+
+# Search CUDA libs from installed python packages.
+if(WIN32)
+  # Resolve paths of unfound DLL at runtime.
+  if(BUILD_SHARED_LIBS)
+    target_link_libraries(mlx PRIVATE "delayimp.lib")
+    target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/delayload.cpp)
+  else()
+    # For static library the delayload must be compiled into final executables.
+    target_link_libraries(mlx PUBLIC "delayimp.lib")
+    target_sources(
+      mlx PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/delayload.cpp>)
+  endif()
+  # Get all the CUDA DLLs we could link with.
+  file(
+    GLOB CUDA_DLL_NAMES
+    RELATIVE "${CUDAToolkit_BIN_DIR}/x64"
+    "${CUDAToolkit_BIN_DIR}/x64/*.dll")
+  # Delay load CUDA and cuDNN libs.
+  foreach(CUDA_DLL ${CUDA_DLL_NAMES} ${CUDNN_DLL_NAMES})
+    target_link_options(mlx PUBLIC "/DELAYLOAD:${CUDA_DLL}")
+  endforeach()
+  # Pass the locations where CUDA DLLs are placed.
+  if(NOT MLX_LOAD_CUDA_LIBS_FROM_PYTHON)
+    target_compile_definitions(
+      mlx PUBLIC MLX_CUDA_BIN_DIR="${CUDAToolkit_BIN_DIR}/x64"
+                 MLX_CUDNN_BIN_DIR="${CUDNN_BIN_DIR}")
+  endif()
+else()
+  # For POSIX we rely on RPATH to search for CUDA libs.
+  if(MLX_LOAD_CUDA_LIBS_FROM_PYTHON)
+    set_property(
+      TARGET mlx
+      APPEND
+      PROPERTY INSTALL_RPATH
+               # The paths here should match the install_requires in setup.py.
+               "$ORIGIN/../../nvidia/cublas/lib"
+               "$ORIGIN/../../nvidia/cuda_nvrtc/lib"
+               "$ORIGIN/../../nvidia/cudnn/lib"
+               "$ORIGIN/../../nvidia/nccl/lib")
+  endif()
 endif()

 # ------------------------ Dependencies ------------------------
@@ -174,7 +214,7 @@ endif()
 # Use fixed version of CCCL.
 FetchContent_Declare(
  cccl
-  URL "https://github.com/NVIDIA/cccl/releases/download/v2.8.1/cccl-v2.8.1.zip")
+  URL "https://github.com/NVIDIA/cccl/releases/download/v3.1.3/cccl-v3.1.3.zip")
 FetchContent_MakeAvailable(cccl)
 target_include_directories(mlx BEFORE PRIVATE "${cccl_SOURCE_DIR}/include")

@@ -202,7 +242,6 @@ FetchContent_MakeAvailable(nvtx3)
 target_link_libraries(mlx PUBLIC $<BUILD_INTERFACE:nvtx3-cpp>)

 # Make cuda runtime APIs available in non-cuda files.
-find_package(CUDAToolkit REQUIRED)
 target_include_directories(mlx PRIVATE ${CUDAToolkit_INCLUDE_DIRS})

 # Use cublasLt.
@@ -225,16 +264,15 @@ set(CUDNN_FRONTEND_BUILD_PYTHON_BINDINGS OFF)
 FetchContent_MakeAvailable(cudnn)
 target_link_libraries(mlx PRIVATE cudnn_frontend)
 # Link with the actual cuDNN libraries.
-include(${cudnn_frontend_SOURCE_DIR}/cmake/cuDNN.cmake)
 target_link_libraries(mlx PRIVATE CUDNN::cudnn_all)

 # Use header-only CUTLASS.
 FetchContent_Declare(
  cutlass
  GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
-  GIT_TAG v4.3.2
+  GIT_TAG v4.3.5
  GIT_SHALLOW TRUE
  SOURCE_SUBDIR include EXCLUDE_FROM_ALL)
 FetchContent_MakeAvailable(cutlass)
 target_include_directories(
-  mlx PRIVATE $<BUILD_INTERFACE:${cutlass_SOURCE_DIR}/include>)
+  mlx SYSTEM PRIVATE $<BUILD_INTERFACE:${cutlass_SOURCE_DIR}/include>)
@@ -3,13 +3,17 @@
 #include "mlx/backend/cuda/allocator.h"
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/utils.h"
+#include "mlx/backend/gpu/device_info.h"
+#include "mlx/memory.h"
+#include "mlx/scheduler.h"
 #include "mlx/utils.h"

 #include <cuda_runtime.h>
 #include <fmt/format.h>
-#include <unistd.h>

 #include <cassert>
+#include <fstream>
+#include <string>

 namespace mlx::core {

@@ -20,6 +24,70 @@ constexpr int page_size = 16384;
 // Any allocations smaller than this will try to use the small pool
 constexpr int small_block_size = 8;

+// The small pool size in bytes. This should be a multiple of the host page
+// size and small_block_size.
+constexpr int small_pool_size = 4 * page_size;
+
+// Check if running on Windows or Windows Subsystem for Linux
+bool is_windows() {
+#if defined(_WIN32)
+  return true;
+#elif defined(__linux__)
+  // WSL kernels contain "microsoft" or "WSL" in /proc/version
+  static bool is_wsl = []() {
+    std::ifstream version("/proc/version");
+    if (version.is_open()) {
+      std::string line;
+      std::getline(version, line);
+      return line.find("microsoft") != std::string::npos ||
+          line.find("Microsoft") != std::string::npos ||
+          line.find("WSL") != std::string::npos;
+    }
+    return false;
+  }();
+  return is_wsl;
+#else
+  return false;
+#endif
+}
+
+bool supports_managed_memory() {
+  static bool managed_memory = []() {
+    int device_count = gpu::device_count();
+    for (int i = 0; i < device_count; ++i) {
+      auto& d = cu::device(i);
+      if (!d.managed_memory()) {
+        return false;
+      }
+      // Empirically on Windows (and WSL) if there is no concurrentManagedAccess
+      // the managed memory also does not work.
+      if (is_windows() && !d.concurrent_managed_access()) {
+        return false;
+      }
+    }
+    return true;
+  }();
+  return managed_memory;
+}
+
+inline void* unified_malloc(size_t size) {
+  void* data = nullptr;
+  if (supports_managed_memory()) {
+    CHECK_CUDA_ERROR(cudaMallocManaged(&data, size));
+  } else {
+    CHECK_CUDA_ERROR(cudaMallocHost(&data, size));
+  }
+  return data;
+}
+
+inline void unified_free(void* data) {
+  if (supports_managed_memory()) {
+    CHECK_CUDA_ERROR(cudaFree(data));
+  } else {
+    CHECK_CUDA_ERROR(cudaFreeHost(data));
+  }
+}
+
 #if CUDART_VERSION >= 13000
 inline cudaMemLocation cuda_mem_loc(int i) {
  cudaMemLocation loc;
@@ -33,24 +101,21 @@ inline int cuda_mem_loc(int i) {
 }
 #endif // CUDART_VERSION >= 13000

-// The small pool size in bytes. This should be a multiple of the host page
-// size and small_block_size.
-constexpr int small_pool_size = 4 * page_size;
-
 SmallSizePool::SmallSizePool() {
  auto num_blocks = small_pool_size / small_block_size;
  buffer_ = new Block[num_blocks];
-
  next_free_ = buffer_;

-  CHECK_CUDA_ERROR(cudaMallocManaged(&data_, small_pool_size));
-
-  int device_count = 0;
-  CHECK_CUDA_ERROR(cudaGetDeviceCount(&device_count));
-  for (int i = 0; i < device_count; ++i) {
-    auto loc = cuda_mem_loc(i);
-    CHECK_CUDA_ERROR(
-        cudaMemAdvise(data_, small_pool_size, cudaMemAdviseSetAccessedBy, loc));
+  data_ = unified_malloc(small_pool_size);
+  if (supports_managed_memory()) {
+    int device_count = gpu::device_count();
+    for (int i = 0; i < device_count; ++i) {
+      if (device(i).concurrent_managed_access()) {
+        auto loc = cuda_mem_loc(i);
+        CHECK_CUDA_ERROR(cudaMemAdvise(
+            data_, small_pool_size, cudaMemAdviseSetAccessedBy, loc));
+      }
+    }
  }

  auto curr = next_free_;
@@ -62,7 +127,7 @@ SmallSizePool::SmallSizePool() {
 }

 SmallSizePool::~SmallSizePool() {
-  CHECK_CUDA_ERROR(cudaFree(data_));
+  unified_free(data_);
  delete[] buffer_;
 }

@@ -96,39 +161,23 @@ CudaAllocator::CudaAllocator()
    : buffer_cache_(
          page_size,
          [](CudaBuffer* buf) { return buf->size; },
-          [this](CudaBuffer* buf) { cuda_free(buf); }) {
+          [this](CudaBuffer* buf) { free_cuda_buffer(buf); }) {
  size_t free;
  CHECK_CUDA_ERROR(cudaMemGetInfo(&free, &total_memory_));
  memory_limit_ = total_memory_ * 0.95;
  free_limit_ = total_memory_ - memory_limit_;
  max_pool_size_ = memory_limit_;

-  int device_count = 0;
-  CHECK_CUDA_ERROR(cudaGetDeviceCount(&device_count));
-  int curr;
-  CHECK_CUDA_ERROR(cudaGetDevice(&curr));
+  int device_count = gpu::device_count();
+  free_streams_.resize(device_count);
+  mem_pools_.resize(device_count);
  for (int i = 0; i < device_count; ++i) {
-    CHECK_CUDA_ERROR(cudaSetDevice(i));
-    cudaStream_t s;
-    CHECK_CUDA_ERROR(cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking));
-    free_streams_.push_back(s);
-
-    cudaMemPool_t mem_pool;
-    CHECK_CUDA_ERROR(cudaDeviceGetDefaultMemPool(&mem_pool, i));
-    mem_pools_.push_back(mem_pool);
+    auto& d = device(i);
+    if (d.memory_pools()) {
+      free_streams_[i] = CudaStream(d);
+      CHECK_CUDA_ERROR(cudaDeviceGetDefaultMemPool(&mem_pools_[i], i));
+    }
  }
-  CHECK_CUDA_ERROR(cudaSetDevice(curr));
-}
-
-void copy_to_managed(CudaBuffer& buf) {
-  // TODO maybe make this async on a i/o stream to avoid synchronizing the
-  // device on malloc/and free
-  void* new_data;
-  CHECK_CUDA_ERROR(cudaMallocManaged(&new_data, buf.size));
-  buf.device = -1;
-  CHECK_CUDA_ERROR(cudaMemcpy(new_data, buf.data, buf.size, cudaMemcpyDefault));
-  CHECK_CUDA_ERROR(cudaFree(buf.data));
-  buf.data = new_data;
 }

 Buffer
@@ -137,8 +186,6 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
    return Buffer{new CudaBuffer{nullptr, 0, -1}};
  }

-  // Find available buffer from cache.
-  std::unique_lock lock(mutex_);
  if (size <= small_block_size) {
    size = 8;
  } else if (size < page_size) {
@@ -151,6 +198,8 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
    device = -1;
  }

+  // Find available buffer from cache.
+  std::unique_lock lock(mutex_);
  CudaBuffer* buf = buffer_cache_.reuse_from_cache(size);
  if (!buf) {
    // If we have a lot of memory pressure try to reclaim memory from the cache.
@@ -168,9 +217,14 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
    if (!buf) {
      void* data = nullptr;
      if (device == -1) {
-        CHECK_CUDA_ERROR(cudaMallocManaged(&data, size));
+        data = unified_malloc(size);
      } else {
-        CHECK_CUDA_ERROR(cudaMallocAsync(&data, size, stream));
+        cu::device(device).make_current();
+        if (mem_pools_[device]) { // supports memory pools
+          CHECK_CUDA_ERROR(cudaMallocAsync(&data, size, stream));
+        } else {
+          CHECK_CUDA_ERROR(cudaMalloc(&data, size));
+        }
      }
      if (!data) {
        std::ostringstream msg;
@@ -186,12 +240,14 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
    // from OOM
    if (get_cache_memory() > 0) {
      for (auto p : mem_pools_) {
-        size_t used = 0;
-        CHECK_CUDA_ERROR(cudaMemPoolGetAttribute(
-            p, cudaMemPoolAttrReservedMemCurrent, &used));
-        if (used > (total_memory_ - free_limit_)) {
-          buffer_cache_.release_cached_buffers(free_limit_);
-          break;
+        if (p) {
+          size_t used = 0;
+          CHECK_CUDA_ERROR(cudaMemPoolGetAttribute(
+              p, cudaMemPoolAttrReservedMemCurrent, &used));
+          if (used > (total_memory_ - free_limit_)) {
+            buffer_cache_.release_cached_buffers(free_limit_);
+            break;
+          }
        }
      }
    }
@@ -203,9 +259,10 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
  if (get_cache_memory() > max_pool_size_) {
    buffer_cache_.release_cached_buffers(get_cache_memory() - max_pool_size_);
  }
-  // Copy to managed here if the buffer is not on the right device
+  lock.unlock();
+  // Copy to unified memory here if the buffer is not on the right device.
  if (buf->device >= 0 && buf->device != device) {
-    copy_to_managed(*buf);
+    move_to_unified_memory(*buf, stream);
  }
  return Buffer{buf};
 }
@@ -229,7 +286,7 @@ void CudaAllocator::free(Buffer buffer) {
  if (get_cache_memory() < max_pool_size_) {
    buffer_cache_.recycle_to_cache(buf);
  } else {
-    cuda_free(buf);
+    free_cuda_buffer(buf);
  }
 }

@@ -241,20 +298,52 @@ size_t CudaAllocator::size(Buffer buffer) const {
  return buf->size;
 }

+void CudaAllocator::move_to_unified_memory(
+    CudaBuffer& buf,
+    cudaStream_t stream) {
+  if (buf.device == -1) {
+    return;
+  }
+  void* data = unified_malloc(buf.size);
+  cudaMemcpyKind kind =
+      supports_managed_memory() ? cudaMemcpyDefault : cudaMemcpyDeviceToHost;
+  if (stream && mem_pools_[buf.device]) {
+    CHECK_CUDA_ERROR(cudaMemcpyAsync(data, buf.data, buf.size, kind, stream));
+    free_async(buf, stream);
+  } else {
+    CHECK_CUDA_ERROR(cudaMemcpy(data, buf.data, buf.size, kind));
+    free_async(buf);
+  }
+  buf.data = data;
+  buf.device = -1;
+}
+
 // This must be called with mutex_ aquired
-void CudaAllocator::cuda_free(CudaBuffer* buf) {
+void CudaAllocator::free_cuda_buffer(CudaBuffer* buf) {
  if (scalar_pool_.in_pool(buf)) {
    scalar_pool_.free(buf);
  } else {
-    if (buf->device >= 0) {
-      CHECK_CUDA_ERROR(cudaFreeAsync(buf->data, free_streams_[buf->device]));
-    } else {
-      CHECK_CUDA_ERROR(cudaFree(buf->data));
-    }
+    free_async(*buf);
    delete buf;
  }
 }

+void CudaAllocator::free_async(CudaBuffer& buf, cudaStream_t stream) {
+  if (buf.device == -1) {
+    unified_free(buf.data);
+  } else {
+    // Free asynchronously when memory pools is supported.
+    if (mem_pools_[buf.device]) {
+      if (!stream) {
+        stream = free_streams_[buf.device];
+      }
+      CHECK_CUDA_ERROR(cudaFreeAsync(buf.data, stream));
+    } else {
+      CHECK_CUDA_ERROR(cudaFree(buf.data));
+    }
+  }
+}
+
 size_t CudaAllocator::get_active_memory() const {
  return active_memory_;
 }
@@ -294,22 +383,20 @@ void CudaAllocator::clear_cache() {
 }

 CudaAllocator& allocator() {
-  // By creating the |allocator_| on heap, the destructor of CudaAllocator
-  // will not be called on exit and buffers in the cache will be leaked. This
-  // can save some time at program exit.
-  static CudaAllocator* allocator_ = new CudaAllocator;
+  static auto* allocator_ = []() {
+    // Ensure scheduler is created before allocator.
+    scheduler::scheduler();
+    // By creating the |allocator_| on heap, the destructor of CudaAllocator
+    // will not be called on exit and buffers in the cache will be leaked. This
+    // can save some time at program exit.
+    return new CudaAllocator();
+  }();
  return *allocator_;
 }

 Buffer malloc_async(size_t size, CommandEncoder& encoder) {
-  auto buffer = allocator().malloc_async(
+  return allocator().malloc_async(
      size, encoder.device().cuda_device(), encoder.stream());
-  if (size && !buffer.ptr()) {
-    std::ostringstream msg;
-    msg << "[malloc_async] Unable to allocate " << size << " bytes.";
-    throw std::runtime_error(msg.str());
-  }
-  return buffer;
 }

 } // namespace cu
@@ -325,9 +412,7 @@ void* Buffer::raw_ptr() {
    return nullptr;
  }
  auto& cbuf = *static_cast<cu::CudaBuffer*>(ptr_);
-  if (cbuf.device != -1) {
-    copy_to_managed(cbuf);
-  }
+  cu::allocator().move_to_unified_memory(cbuf);
  return cbuf.data;
 }

@@ -54,6 +54,10 @@ class CudaAllocator : public allocator::Allocator {
  void free(Buffer buffer) override;
  size_t size(Buffer buffer) const override;

+  // Replace the memory of |buf| with unified memory (managed memory or pinned
+  // host memory), and copy the data over. Pass |stream| to copy asynchronously.
+  void move_to_unified_memory(CudaBuffer& buf, cudaStream_t stream = nullptr);
+
  size_t get_active_memory() const;
  size_t get_peak_memory() const;
  void reset_peak_memory();
@@ -64,7 +68,8 @@ class CudaAllocator : public allocator::Allocator {
  void clear_cache();

 private:
-  void cuda_free(CudaBuffer* buf);
+  void free_cuda_buffer(CudaBuffer* buf);
+  void free_async(CudaBuffer& buf, cudaStream_t stream = nullptr);

  CudaAllocator();
  friend CudaAllocator& allocator();
@@ -77,7 +82,7 @@ class CudaAllocator : public allocator::Allocator {
  BufferCache<CudaBuffer> buffer_cache_;
  size_t active_memory_{0};
  size_t peak_memory_{0};
-  std::vector<cudaStream_t> free_streams_;
+  std::vector<CudaStream> free_streams_;
  std::vector<cudaMemPool_t> mem_pools_;
  SmallSizePool scalar_pool_;
 };
@@ -56,7 +56,6 @@ void Arange::eval_gpu(const std::vector<array>& inputs, array& out) {
          cu::arange<OutType, IdxT, N_WRITES>,
          num_blocks,
          block_dims,
-          0,
          gpu_ptr<OutType>(out),
          out.data_size(),
          static_cast<CTYPE>(start_),
@@ -172,7 +172,6 @@ void ArgReduce::eval_gpu(const std::vector<array>& inputs, array& out) {
          kernel,
          num_blocks,
          block_dim(),
-          0,
          gpu_ptr<T>(in),
          gpu_ptr<uint32_t>(out),
          out.size(),
@@ -291,7 +291,6 @@ void binary_op_gpu_inplace(
                        kernel,
                        {num_blocks_x, num_blocks_y},
                        block_dims,
-                        0,
                        gpu_ptr<InType>(a),
                        gpu_ptr<InType>(b),
                        gpu_ptr<OutType>(out),
@@ -309,7 +308,6 @@ void binary_op_gpu_inplace(
                      kernel,
                      {num_blocks_x, num_blocks_y},
                      block_dims,
-                      0,
                      gpu_ptr<InType>(a),
                      gpu_ptr<InType>(b),
                      gpu_ptr<OutType>(out),
@@ -338,7 +336,6 @@ void binary_op_gpu_inplace(
                kernel,
                num_blocks,
                block_dims,
-                0,
                gpu_ptr<InType>(a),
                gpu_ptr<InType>(b),
                gpu_ptr<OutType>(out),
@@ -346,11 +343,12 @@ void binary_op_gpu_inplace(
          });
        }
      } else {
-        throw std::runtime_error(fmt::format(
-            "Can not do binary op {} on inputs of {} with result of {}.",
-            op,
-            dtype_to_string(a.dtype()),
-            dtype_to_string(out.dtype())));
+        throw std::runtime_error(
+            fmt::format(
+                "Can not do binary op {} on inputs of {} with result of {}.",
+                op,
+                dtype_to_string(a.dtype()),
+                dtype_to_string(out.dtype())));
      }
    });
  });
@@ -314,7 +314,6 @@ void binary_two_op_gpu_inplace(
                        kernel,
                        {num_blocks_x, num_blocks_y},
                        block_dims,
-                        0,
                        gpu_ptr<InType>(a),
                        gpu_ptr<InType>(b),
                        gpu_ptr<OutType>(out_a),
@@ -333,7 +332,6 @@ void binary_two_op_gpu_inplace(
                      kernel,
                      {num_blocks_x, num_blocks_y},
                      block_dims,
-                      0,
                      gpu_ptr<InType>(a),
                      gpu_ptr<InType>(b),
                      gpu_ptr<OutType>(out_a),
@@ -367,7 +365,6 @@ void binary_two_op_gpu_inplace(
                kernel,
                num_blocks,
                block_dims,
-                0,
                gpu_ptr<InType>(a),
                gpu_ptr<InType>(b),
                gpu_ptr<OutType>(out_a),
@@ -376,11 +373,12 @@ void binary_two_op_gpu_inplace(
          });
        }
      } else {
-        throw std::runtime_error(fmt::format(
-            "Can not do binary op {} on inputs of {} with result of {}.",
-            op,
-            dtype_to_string(a.dtype()),
-            dtype_to_string(out_a.dtype())));
+        throw std::runtime_error(
+            fmt::format(
+                "Can not do binary op {} on inputs of {} with result of {}.",
+                op,
+                dtype_to_string(a.dtype()),
+                dtype_to_string(out_a.dtype())));
      }
    });
  });
@@ -36,14 +36,16 @@ struct FusedKernelBuilder {
      params.push_back(
          fmt::format("const {}* {}", dtype_to_cuda_type(x.dtype()), xname));
      if (!is_scalar(x) && !contiguous) {
-        params.push_back(fmt::format(
-            "const __grid_constant__ cuda::std::array<int64_t, NDIM> {}_strides",
-            xname));
+        params.push_back(
+            fmt::format(
+                "const __grid_constant__ cuda::std::array<int64_t, NDIM> {}_strides",
+                xname));
      }
    }
    for (const auto& x : outputs) {
-      params.push_back(fmt::format(
-          "{}* {}", dtype_to_cuda_type(x.dtype()), namer.get_name(x)));
+      params.push_back(
+          fmt::format(
+              "{}* {}", dtype_to_cuda_type(x.dtype()), namer.get_name(x)));
    }
    if (!contiguous) {
      params.push_back(
@@ -250,20 +252,30 @@ void Compiled::eval_gpu(
    builder.os += "\n} // namespace mlx::core::cu\n";
    // Build kernel names.
    std::vector<std::string> kernel_names;
-    kernel_names.push_back(fmt::format(
-        "mlx::core::cu::{}_contiguous<uint32_t, {}>",
-        lib_name(),
-        work_per_thread));
-    kernel_names.push_back(fmt::format(
-        "mlx::core::cu::{}_contiguous<int64_t, {}>",
-        lib_name(),
-        work_per_thread));
-    for (auto wpt : std::array<int, 2>{1, work_per_thread}) {
+    kernel_names.push_back(
+        fmt::format(
+            "mlx::core::cu::{}_contiguous<uint32_t, {}>",
+            lib_name(),
+            work_per_thread));
+    kernel_names.push_back(
+        fmt::format(
+            "mlx::core::cu::{}_contiguous<int64_t, {}>",
+            lib_name(),
+            work_per_thread));
+    for (int wpt : {1, work_per_thread}) {
      for (int i = 1; i <= MAX_NDIM; ++i) {
-        kernel_names.push_back(fmt::format(
-            "mlx::core::cu::{}_strided<{}, uint32_t, {}>", lib_name(), i, wpt));
-        kernel_names.push_back(fmt::format(
-            "mlx::core::cu::{}_strided<{}, int64_t, {}>", lib_name(), i, wpt));
+        kernel_names.push_back(
+            fmt::format(
+                "mlx::core::cu::{}_strided<{}, uint32_t, {}>",
+                lib_name(),
+                i,
+                wpt));
+        kernel_names.push_back(
+            fmt::format(
+                "mlx::core::cu::{}_strided<{}, int64_t, {}>",
+                lib_name(),
+                i,
+                wpt));
      }
    }

@@ -339,7 +351,8 @@ void Compiled::eval_gpu(
  auto [kernel, max_block_dims] = mod.get_kernel_and_dims(kernel_name);
  auto [num_blocks, block_dims] =
      get_launch_args(outputs[0], large, work_per_thread, max_block_dims);
-  encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
+  encoder.add_kernel_node_raw(
+      kernel, num_blocks, block_dims, {}, 0, args.args());
 }

 } // namespace mlx::core
@@ -103,7 +103,7 @@ std::optional<DnnGraph> build_conv_graph(
    const std::vector<int64_t>& dilation) {
  auto compute_dtype =
      (dtype == float16 || dtype == bfloat16) ? float32 : dtype;
-  DnnGraph graph(encoder.device().cudnn_handle(), dtype, compute_dtype);
+  DnnGraph graph(encoder.device().get_cudnn_handle(), dtype, compute_dtype);
  auto x_ = graph.tensor_nchw("X", 'x', x);
  auto w_ = graph.tensor_nchw("W", 'w', w);

@@ -117,7 +117,6 @@ array unfold_inputs_nd(
        cu::naive_unfold_nd<DataType, NDIM>,
        num_blocks,
        block_dims,
-        0,
        gpu_ptr<DataType>(in),
        gpu_ptr<DataType>(unfolded),
        filter_size,
@@ -120,7 +120,6 @@ array grouped_unfold_transpose_inputs_nd(
        cu::naive_grouped_unfold_transpose_nd<DataType, NDIM>,
        num_blocks,
        block_dims,
-        0,
        gpu_ptr<DataType>(in),
        gpu_ptr<DataType>(unfolded),
        filter_size,
@@ -76,7 +76,6 @@ void copy_contiguous(
            kernel,
            num_blocks,
            block_dims,
-            0,
            gpu_ptr<InType>(in) + in_offset,
            gpu_ptr<OutType>(out) + out_offset,
            out.data_size());
@@ -137,7 +137,6 @@ void copy_general(
                    kernel,
                    {num_blocks_x, num_blocks_y},
                    block_dims,
-                    0,
                    in_ptr,
                    out_ptr,
                    rest,
@@ -154,7 +153,6 @@ void copy_general(
                  kernel,
                  {num_blocks_x, num_blocks_y},
                  block_dims,
-                  0,
                  in_ptr,
                  out_ptr,
                  rest,
@@ -83,7 +83,6 @@ void copy_general_dynamic(
                        dims_constant()>,
                    num_blocks,
                    block_dims,
-                    0,
                    in_ptr,
                    out_ptr,
                    out.size(),
@@ -99,7 +98,6 @@ void copy_general_dynamic(
                  cu::copy_gg_dynamic<InType, OutType, IdxT>,
                  num_blocks,
                  block_dims,
-                  0,
                  in_ptr,
                  out_ptr,
                  out.size(),
@@ -154,7 +154,6 @@ void copy_general_input(
            kernel,
            {num_blocks_x, num_blocks_y},
            block_dims,
-            0,
            in_ptr,
            out_ptr,
            int64_t(shape[0]),
@@ -195,7 +194,6 @@ void copy_general_input(
                    kernel,
                    {num_blocks_x, num_blocks_y},
                    block_dims,
-                    0,
                    in_ptr,
                    out_ptr,
                    rest,
@@ -213,7 +211,6 @@ void copy_general_input(
                  kernel,
                  {num_blocks_x, num_blocks_y},
                  block_dims,
-                  0,
                  in_ptr,
                  out_ptr,
                  rest,
@@ -98,20 +98,13 @@ void CublasMatmulBase::init_base(
  M_ = a_rows;
  N_ = b_cols;
  scale_type_ = scale_type;
-  handle_ = device.lt_handle();
+  handle_ = device.get_cublaslt_handle();
  pref_ = cublas_utils::get_preference(device);
  heuristic_.state = CUBLAS_STATUS_NOT_INITIALIZED;

  CHECK_CUBLAS_ERROR(
      cublasLtMatmulDescCreate(&matmul_desc_, compute_type, scale_type));

-  int32_t pointer_mode = CUBLASLT_POINTER_MODE_HOST;
-  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
-      matmul_desc_,
-      CUBLASLT_MATMUL_DESC_POINTER_MODE,
-      &pointer_mode,
-      sizeof(int32_t)));
-
  // In cublasLt matrices use column-major layout, while it is possible to use
  // the CUBLASLT_ORDER_ROW option to switch to row-major layout, the bias
  // epilogue does not work with the option. So instead we swap A and B to make
@@ -34,8 +34,9 @@ inline cudaDataType_t dtype_to_cublas_type(Dtype dtype, std::string_view tag) {
    case complex64:
      return CUDA_C_32F;
    default:
-      throw std::runtime_error(fmt::format(
-          "Unsupported dtype in {}: {}.", tag, dtype_to_string(dtype)));
+      throw std::runtime_error(
+          fmt::format(
+              "Unsupported dtype in {}: {}.", tag, dtype_to_string(dtype)));
  }
 }

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

 #pragma once

+#include <string>
+#include <unordered_map>
+#include <variant>
+
+#include "mlx/api.h"
+
 namespace mlx::core::cu {

 /* Check if the CUDA backend is available. */
-bool is_available();
+MLX_API bool is_available();
+
+/* Get information about a CUDA device. */
+MLX_API const
+    std::unordered_map<std::string, std::variant<std::string, size_t>>&
+    device_info(int device_index = 0);

 } // namespace mlx::core::cu
@@ -83,6 +83,7 @@ class CudaGraphExec : public CudaHandle<cudaGraphExec_t, cudaGraphExecDestroy> {

 class CudaStream : public CudaHandle<cudaStream_t, cudaStreamDestroy> {
 public:
+  using CudaHandle::CudaHandle;
  explicit CudaStream(cu::Device& device);
 };

@@ -62,8 +62,9 @@ inline fe::DataType_t dtype_to_cudnn_type(Dtype dtype) {
    case float64:
      return fe::DataType_t::DOUBLE;
    default:
-      throw std::runtime_error(fmt::format(
-          "Unsupported dtype in cuDNN: {}.", dtype_to_string(dtype)));
+      throw std::runtime_error(
+          fmt::format(
+              "Unsupported dtype in cuDNN: {}.", dtype_to_string(dtype)));
  }
 }

@@ -72,13 +73,13 @@ inline fe::DataType_t dtype_to_cudnn_type(Dtype dtype) {
 // There are 2 differences from the const_param util from kernel_utils.cuh:
 // 1. The rest of array is filled with 0.
 // 2. This util can be used in .cpp files.
-template <int NDIM = MAX_NDIM, typename T, template <typename U> class Vec>
-inline std::array<T, NDIM> vector_key(const Vec<T>& vec) {
+template <int NDIM = MAX_NDIM, typename Vec>
+inline std::array<typename Vec::value_type, NDIM> vector_key(const Vec& vec) {
  if (vec.size() > NDIM) {
    throw std::runtime_error(
        fmt::format("ndim can not be larger than {}.", NDIM));
  }
-  std::array<T, NDIM> result = {};
+  std::array<typename Vec::value_type, NDIM> result = {};
  std::copy_n(vec.begin(), vec.size(), result.begin());
  return result;
 }
@@ -122,15 +123,30 @@ class DnnGraph : public fe::graph::Graph {
    return attrs;
  }

+  // Create a 4D cuDNN tensor from 1D array, with |axis| being contiguous dim.
+  auto tensor_4d(const char* name, int64_t uid, const array& x, int axis) {
+    assert(x.ndim() == 1);
+    auto attrs = Graph::tensor(fe::graph::Tensor_attributes().set_name(name));
+    std::vector<int64_t> shape(4, 1);
+    std::vector<int64_t> strides(4, 1);
+    shape.at(axis) = x.size();
+    if (axis > 0) {
+      strides.at(axis - 1) = x.size();
+    }
+    set_tensor_attrs(attrs, uid, x, shape, strides);
+    return attrs;
+  }
+
  // Create a cuDNN tensor for scalar.
  auto scalar(const char* name, int64_t uid, Dtype dtype) {
-    return Graph::tensor(fe::graph::Tensor_attributes()
-                             .set_name(name)
-                             .set_uid(uid)
-                             .set_dim({1, 1, 1, 1})
-                             .set_stride({1, 1, 1, 1})
-                             .set_is_pass_by_value(true)
-                             .set_data_type(dtype_to_cudnn_type(dtype)));
+    return Graph::tensor(
+        fe::graph::Tensor_attributes()
+            .set_name(name)
+            .set_uid(uid)
+            .set_dim({1, 1, 1, 1})
+            .set_stride({1, 1, 1, 1})
+            .set_is_pass_by_value(true)
+            .set_data_type(dtype_to_cudnn_type(dtype)));
  }

  // Call this before setting notes.
@@ -373,7 +373,8 @@ void CustomKernel::eval_gpu(
              kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem);
        }
      });
-  encoder.add_kernel_node(kernel, grid, block, shared_memory_, args.args());
+  encoder.add_kernel_node_raw(
+      kernel, grid, block, {}, shared_memory_, args.args());
 }

 } // namespace mlx::core::fast
@@ -0,0 +1,46 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/dtype.h"
+
+#include <cutlass/bfloat16.h>
+#include <cutlass/half.h>
+#include <fmt/format.h>
+
+namespace mlx::core {
+
+// Throw exception if the cutlass API does not succeed.
+inline void check_cutlass_error(const char* name, cutlass::Status status) {
+  if (status != cutlass::Status::kSuccess) {
+    throw std::runtime_error(
+        fmt::format(
+            "{} failed with code: {}.",
+            name,
+            cutlass::cutlassGetStatusString(status)));
+  }
+}
+
+// The macro version that prints the command that failed.
+#define CHECK_CUTLASS_ERROR(cmd) ::mlx::core::check_cutlass_error(#cmd, (cmd))
+
+// Maps CPU types to CUTLASS types.
+template <typename T>
+struct CTypeToCutlassType {
+  using type = T;
+};
+
+template <>
+struct CTypeToCutlassType<float16_t> {
+  using type = cutlass::half_t;
+};
+
+template <>
+struct CTypeToCutlassType<bfloat16_t> {
+  using type = cutlass::bfloat16_t;
+};
+
+template <typename T>
+using cutlass_type_t = typename CTypeToCutlassType<T>::type;
+
+} // namespace mlx::core
@@ -0,0 +1,80 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/common/utils.h"
+
+// clang-format off
+#include <windows.h> // must be included first
+#include <delayimp.h>
+// clang-format on
+
+namespace mlx::core {
+
+namespace fs = std::filesystem;
+
+inline fs::path relative_to_current_binary(const char* relative) {
+  return fs::absolute(current_binary_dir() / relative);
+}
+
+inline fs::path cublas_bin_dir() {
+#if defined(MLX_CUDA_BIN_DIR)
+  return MLX_CUDA_BIN_DIR;
+#else
+  return relative_to_current_binary("../nvidia/cublas/bin");
+#endif
+}
+
+fs::path load_nvrtc() {
+#if defined(MLX_CUDA_BIN_DIR)
+  fs::path nvrtc_bin_dir = MLX_CUDA_BIN_DIR;
+#else
+  fs::path nvrtc_bin_dir =
+      relative_to_current_binary("../nvidia/cuda_nvrtc/bin");
+#endif
+  // Internally nvrtc loads some libs dynamically, add to search dirs.
+  ::AddDllDirectory(nvrtc_bin_dir.c_str());
+  return nvrtc_bin_dir;
+}
+
+fs::path load_cudnn() {
+#if defined(MLX_CUDNN_BIN_DIR)
+  fs::path cudnn_bin_dir = MLX_CUDNN_BIN_DIR;
+#else
+  fs::path cudnn_bin_dir = relative_to_current_binary("../nvidia/cudnn/bin");
+#endif
+  // Must load cudnn_graph64_9.dll before locating symbols, otherwise We would
+  // get errors like "Invalid handle. Cannot load symbol cudnnCreate".
+  for (const auto& dll : fs::directory_iterator(cudnn_bin_dir)) {
+    if (dll.path().filename().string().starts_with("cudnn_graph") &&
+        dll.path().extension() == ".dll") {
+      ::LoadLibraryW(dll.path().c_str());
+      break;
+    }
+  }
+  // Internally cuDNN loads some libs dynamically, add to search dirs.
+  load_nvrtc();
+  ::AddDllDirectory(cudnn_bin_dir.c_str());
+  ::AddDllDirectory(cublas_bin_dir().c_str());
+  return cudnn_bin_dir;
+}
+
+// Called by system when failed to locate a lazy-loaded DLL.
+FARPROC WINAPI delayload_helper(unsigned dliNotify, PDelayLoadInfo pdli) {
+  HMODULE mod = NULL;
+  if (dliNotify == dliNotePreLoadLibrary) {
+    std::string dll = pdli->szDll;
+    if (dll.starts_with("cudnn")) {
+      static auto cudnn_bin_dir = load_cudnn();
+      mod = ::LoadLibraryW((cudnn_bin_dir / dll).c_str());
+    } else if (dll.starts_with("cublas")) {
+      mod = ::LoadLibraryW((cublas_bin_dir() / dll).c_str());
+    } else if (dll.starts_with("nvrtc")) {
+      static auto nvrtc_bin_dir = load_nvrtc();
+      mod = ::LoadLibraryW((nvrtc_bin_dir / dll).c_str());
+    }
+  }
+  return reinterpret_cast<FARPROC>(mod);
+}
+
+} // namespace mlx::core
+
+extern "C" const PfnDliHook __pfnDliNotifyHook2 = mlx::core::delayload_helper;
@@ -3,6 +3,7 @@
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/jit_module.h"
 #include "mlx/backend/cuda/worker.h"
+#include "mlx/backend/gpu/device_info.h"
 #include "mlx/utils.h"

 #include <fmt/format.h>
@@ -30,6 +31,11 @@ const char* save_cuda_graphs_dot_file() {
  return filename;
 }

+inline bool is_empty_dim(dim3 dim) {
+  return (dim.x == 0 && dim.y == 0 && dim.z == 0) ||
+      (dim.x == 1 && dim.y == 1 && dim.z == 1);
+}
+
 } // namespace

 Device::Device(int device) : device_(device) {
@@ -37,36 +43,33 @@ Device::Device(int device) : device_(device) {
      &compute_capability_major_, cudaDevAttrComputeCapabilityMajor, device_));
  CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
      &compute_capability_minor_, cudaDevAttrComputeCapabilityMinor, device_));
-  // Validate the requirements of device.
-  int attr = 0;
  CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
-      &attr, cudaDevAttrConcurrentManagedAccess, device_));
-  if (attr != 1) {
-    throw std::runtime_error(fmt::format(
-        "Device {} does not support synchronization in managed memory.",
-        device_));
-  }
-
-  // The cublasLt handle is used by matmul.
-  make_current();
-  CHECK_CUBLAS_ERROR(cublasLtCreate(&lt_));
-  // The cudnn handle is used by Convolution.
-  CHECK_CUDNN_ERROR(cudnnCreate(&cudnn_));
-
-  // Initialize the jit module cache here ensures it is not
-  // unloaded before any evaluation is done
-  get_jit_module_cache();
+      &concurrent_managed_access_,
+      cudaDevAttrConcurrentManagedAccess,
+      device_));
+  CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
+      &host_native_atomic_, cudaDevAttrHostNativeAtomicSupported, device_));
+  CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
+      &managed_memory_, cudaDevAttrManagedMemory, device_));
+  CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
+      &memory_pools_, cudaDevAttrMemoryPoolsSupported, device_));
 }

 Device::~Device() {
-  CHECK_CUDNN_ERROR(cudnnDestroy(cudnn_));
-  CHECK_CUBLAS_ERROR(cublasLtDestroy(lt_));
+  if (cudnn_handle_) {
+    CHECK_CUDNN_ERROR(cudnnDestroy(cudnn_handle_));
+  }
+  if (cublaslt_handle_) {
+    CHECK_CUBLAS_ERROR(cublasLtDestroy(cublaslt_handle_));
+  }
 }

 void Device::make_current() {
  // We need to set/get current CUDA device very frequently, cache it to reduce
-  // actual calls of CUDA APIs.
-  static thread_local int current = 0;
+  // actual calls of CUDA APIs. Use -1 as sentinel so the first call on each
+  // new thread always calls cudaSetDevice (which establishes the CUDA primary
+  // context). Without this, device 0 would never get set on a new thread.
+  static thread_local int current = -1;
  if (current != device_) {
    CHECK_CUDA_ERROR(cudaSetDevice(device_));
    current = device_;
@@ -81,6 +84,22 @@ CommandEncoder& Device::get_command_encoder(Stream s) {
  return it->second;
 }

+cublasLtHandle_t Device::get_cublaslt_handle() {
+  if (!cublaslt_handle_) {
+    make_current();
+    CHECK_CUBLAS_ERROR(cublasLtCreate(&cublaslt_handle_));
+  }
+  return cublaslt_handle_;
+}
+
+cudnnHandle_t Device::get_cudnn_handle() {
+  if (!cudnn_handle_) {
+    make_current();
+    CHECK_CUDNN_ERROR(cudnnCreate(&cudnn_handle_));
+  }
+  return cudnn_handle_;
+}
+
 CommandEncoder::CaptureContext::CaptureContext(CommandEncoder& enc) : enc(enc) {
  enc.device().make_current();
  if (!use_cuda_graphs()) {
@@ -202,12 +221,10 @@ std::pair<int, int> get_graph_limits(Device& d) {
      mb = 400;
      break;
    case 900: // H100
-      ops = 30;
-      mb = 400;
-      break;
    case 1000: // B200
-      ops = 50;
-      mb = 500;
+    case 1200: // Consumer Blackwell
+      ops = 100;
+      mb = 1000;
      break;
    case 1210: // DGX Spark
      ops = 20;
@@ -249,51 +266,88 @@ void CommandEncoder::set_output_array(const array& arr) {
  active_outputs_.push_back(id);
 }

-void CommandEncoder::add_kernel_node(
+void CommandEncoder::add_kernel_node_raw(
    void* func,
    dim3 grid_dim,
    dim3 block_dim,
+    dim3 cluster_dim,
    uint32_t smem_bytes,
    void** params) {
+  bool use_cluster = !is_empty_dim(cluster_dim);
+  assert(!use_cluster || device_.compute_capability_major() >= 9);
+
  if (!use_cuda_graphs()) {
    node_count_++;
-    CHECK_CUDA_ERROR(cudaLaunchKernel(
-        func, grid_dim, block_dim, params, smem_bytes, stream()));
+    cudaLaunchConfig_t config = {};
+    config.gridDim = grid_dim;
+    config.blockDim = block_dim;
+    config.dynamicSmemBytes = smem_bytes;
+    config.stream = stream();
+    cudaLaunchAttribute attr = {};
+    if (use_cluster) {
+      attr.id = cudaLaunchAttributeClusterDimension;
+      attr.val.clusterDim.x = cluster_dim.x;
+      attr.val.clusterDim.y = cluster_dim.y;
+      attr.val.clusterDim.z = cluster_dim.z;
+      config.attrs = &attr;
+      config.numAttrs = 1;
+    }
+    CHECK_CUDA_ERROR(cudaLaunchKernelExC(&config, func, params));
    return;
  }
+
  cudaKernelNodeParams kernel_params = {0};
  kernel_params.func = func;
  kernel_params.gridDim = grid_dim;
  kernel_params.blockDim = block_dim;
  kernel_params.kernelParams = params;
  kernel_params.sharedMemBytes = smem_bytes;
-  add_kernel_node(kernel_params);
+  cudaGraphNode_t node = add_kernel_node_raw(kernel_params);
+  if (use_cluster) {
+    cudaKernelNodeAttrValue attr = {};
+    attr.clusterDim.x = cluster_dim.x;
+    attr.clusterDim.y = cluster_dim.y;
+    attr.clusterDim.z = cluster_dim.z;
+    CHECK_CUDA_ERROR(cudaGraphKernelNodeSetAttribute(
+        node, cudaLaunchAttributeClusterDimension, &attr));
+  }
 }

-void CommandEncoder::add_kernel_node(
+void CommandEncoder::add_kernel_node_raw(
    CUfunction func,
    dim3 grid_dim,
    dim3 block_dim,
+    dim3 cluster_dim,
    uint32_t smem_bytes,
    void** params) {
+  bool use_cluster = !is_empty_dim(cluster_dim);
+  assert(!use_cluster || device_.compute_capability_major() >= 9);
+
  if (!use_cuda_graphs()) {
    node_count_++;
-    CHECK_CUDA_ERROR(cuLaunchKernel(
-        func,
-        grid_dim.x,
-        grid_dim.y,
-        grid_dim.z,
-        block_dim.x,
-        block_dim.y,
-        block_dim.z,
-        smem_bytes,
-        stream(),
-        params,
-        nullptr));
+    CUlaunchConfig config = {};
+    config.gridDimX = grid_dim.x;
+    config.gridDimY = grid_dim.y;
+    config.gridDimZ = grid_dim.z;
+    config.blockDimX = block_dim.x;
+    config.blockDimY = block_dim.y;
+    config.blockDimZ = block_dim.z;
+    config.sharedMemBytes = smem_bytes;
+    config.hStream = stream();
+    CUlaunchAttribute attr = {};
+    if (use_cluster) {
+      attr.id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
+      attr.value.clusterDim.x = cluster_dim.x;
+      attr.value.clusterDim.y = cluster_dim.y;
+      attr.value.clusterDim.z = cluster_dim.z;
+      config.attrs = &attr;
+      config.numAttrs = 1;
+    }
+    CHECK_CUDA_ERROR(cuLaunchKernelEx(&config, func, params, nullptr));
    return;
  }

-  CUDA_KERNEL_NODE_PARAMS kernel_params = {0};
+  CUDA_KERNEL_NODE_PARAMS kernel_params = {};
  kernel_params.func = func;
  kernel_params.gridDimX = grid_dim.x;
  kernel_params.gridDimY = grid_dim.y;
@@ -303,19 +357,31 @@ void CommandEncoder::add_kernel_node(
  kernel_params.blockDimZ = block_dim.z;
  kernel_params.kernelParams = params;
  kernel_params.sharedMemBytes = smem_bytes;
-  add_kernel_node(kernel_params);
+  CUgraphNode node = add_kernel_node_raw(kernel_params);
+  if (use_cluster) {
+    CUlaunchAttributeValue attr = {};
+    attr.clusterDim.x = cluster_dim.x;
+    attr.clusterDim.y = cluster_dim.y;
+    attr.clusterDim.z = cluster_dim.z;
+    CHECK_CUDA_ERROR(cuGraphKernelNodeSetAttribute(
+        node, CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION, &attr));
+  }
 }

-void CommandEncoder::add_kernel_node(const cudaKernelNodeParams& params) {
+cudaGraphNode_t CommandEncoder::add_kernel_node_raw(
+    const cudaKernelNodeParams& params) {
  cudaGraphNode_t node;
  CHECK_CUDA_ERROR(cudaGraphAddKernelNode(&node, graph_, NULL, 0, &params));
  insert_graph_dependencies(GraphNode{node, "K"});
+  return node;
 }

-void CommandEncoder::add_kernel_node(const CUDA_KERNEL_NODE_PARAMS& params) {
+CUgraphNode CommandEncoder::add_kernel_node_raw(
+    const CUDA_KERNEL_NODE_PARAMS& params) {
  CUgraphNode node;
  CHECK_CUDA_ERROR(cuGraphAddKernelNode(&node, graph_, NULL, 0, &params));
  insert_graph_dependencies(GraphNode{node, "K"});
+  return node;
 }

 std::pair<std::string, bool> subgraph_to_key(cudaGraph_t graph) {
@@ -409,14 +475,17 @@ void CommandEncoder::commit() {
  }
  if (use_cuda_graphs() && node_count_ > 0) {
    if (!from_nodes_.empty()) {
+#if CUDART_VERSION >= 13000
      CHECK_CUDA_ERROR(cudaGraphAddDependencies(
          graph_,
          from_nodes_.data(),
          to_nodes_.data(),
-#if CUDART_VERSION >= 13000
          nullptr, // edgeData
-#endif // CUDART_VERSION >= 13000
          from_nodes_.size()));
+#else
+      CHECK_CUDA_ERROR(cudaGraphAddDependencies(
+          graph_, from_nodes_.data(), to_nodes_.data(), from_nodes_.size()));
+#endif
    }

    device_.make_current();
@@ -483,13 +552,23 @@ void CommandEncoder::synchronize() {
  f.wait();
 }

-Device& device(mlx::core::Device device) {
-  static std::unordered_map<int, Device> devices;
-  auto it = devices.find(device.index);
-  if (it == devices.end()) {
-    it = devices.try_emplace(device.index, device.index).first;
-  }
-  return it->second;
+Device& device(int cuda_device) {
+  static auto devices = []() {
+    std::vector<Device> devices;
+    int device_count = gpu::device_count();
+    for (int i = 0; i < device_count; ++i) {
+      devices.emplace_back(i);
+    }
+    // Initialize the jit module cache here ensures it is not unloaded before
+    // any evaluation is done.
+    get_jit_module_cache();
+    return devices;
+  }();
+  return devices.at(cuda_device);
+}
+
+Device& device(mlx::core::Device d) {
+  return device(d.index);
 }

 CommandEncoder& get_command_encoder(Stream s) {
@@ -11,7 +11,6 @@
 #include <cublasLt.h>
 #include <cuda.h>
 #include <cudnn.h>
-#include <thrust/execution_policy.h>

 #include <unordered_map>

@@ -48,10 +47,17 @@ class CommandEncoder {
  void set_output_array(const array& arr);

  template <typename F, typename... Params>
-  void add_kernel_node(
+  void
+  add_kernel_node(F* func, dim3 grid_dim, dim3 block_dim, Params&&... params) {
+    add_kernel_node_ex(func, grid_dim, block_dim, {}, 0, params...);
+  }
+
+  template <typename F, typename... Params>
+  void add_kernel_node_ex(
      F* func,
      dim3 grid_dim,
      dim3 block_dim,
+      dim3 cluster_dim,
      uint32_t smem_bytes,
      Params&&... params) {
    constexpr size_t num = sizeof...(Params);
@@ -60,20 +66,28 @@ class CommandEncoder {
    ([&](auto&& p) { ptrs[i++] = static_cast<void*>(&p); }(
         std::forward<Params>(params)),
     ...);
-    add_kernel_node((void*)func, grid_dim, block_dim, smem_bytes, ptrs);
+    add_kernel_node_raw(
+        reinterpret_cast<void*>(func),
+        grid_dim,
+        block_dim,
+        cluster_dim,
+        smem_bytes,
+        ptrs);
  }

-  void add_kernel_node(
-      CUfunction func,
-      dim3 grid_dim,
-      dim3 block_dim,
-      uint32_t smem_bytes,
-      void** params);
-
-  void add_kernel_node(
+  void add_kernel_node_raw(
      void* func,
      dim3 grid_dim,
      dim3 block_dim,
+      dim3 cluster_dim,
+      uint32_t smem_bytes,
+      void** params);
+
+  void add_kernel_node_raw(
+      CUfunction func,
+      dim3 grid_dim,
+      dim3 block_dim,
+      dim3 cluster_dim,
      uint32_t smem_bytes,
      void** params);

@@ -99,8 +113,8 @@ class CommandEncoder {
  void synchronize();

 private:
-  void add_kernel_node(const cudaKernelNodeParams& params);
-  void add_kernel_node(const CUDA_KERNEL_NODE_PARAMS& params);
+  cudaGraphNode_t add_kernel_node_raw(const cudaKernelNodeParams& params);
+  CUgraphNode add_kernel_node_raw(const CUDA_KERNEL_NODE_PARAMS& params);

  struct GraphNode {
    cudaGraphNode_t node;
@@ -119,7 +133,7 @@ class CommandEncoder {
  CudaStream stream_;
  CudaGraph graph_;
  Worker worker_;
-  char node_count_{0};
+  int node_count_{0};
  bool in_concurrent_{false};
  std::vector<cudaGraphNode_t> from_nodes_;
  std::vector<cudaGraphNode_t> to_nodes_;
@@ -142,6 +156,7 @@ class Device {
  explicit Device(int device);
  ~Device();

+  Device(Device&&) = default;
  Device(const Device&) = delete;
  Device& operator=(const Device&) = delete;

@@ -149,6 +164,8 @@ class Device {
  void make_current();

  CommandEncoder& get_command_encoder(Stream s);
+  cublasLtHandle_t get_cublaslt_handle();
+  cudnnHandle_t get_cudnn_handle();

  int cuda_device() const {
    return device_;
@@ -159,31 +176,35 @@ class Device {
  int compute_capability_minor() const {
    return compute_capability_minor_;
  }
-  cublasLtHandle_t lt_handle() const {
-    return lt_;
+  bool concurrent_managed_access() const {
+    return concurrent_managed_access_ == 1;
  }
-  cudnnHandle_t cudnn_handle() const {
-    return cudnn_;
+  bool host_native_atomic() const {
+    return host_native_atomic_ == 1;
+  }
+  bool managed_memory() const {
+    return managed_memory_ == 1;
+  }
+  bool memory_pools() const {
+    return memory_pools_ == 1;
  }

 private:
  int device_;
  int compute_capability_major_;
  int compute_capability_minor_;
+  int concurrent_managed_access_;
+  int host_native_atomic_;
+  int managed_memory_;
+  int memory_pools_;
  std::string device_name_;
-  cublasLtHandle_t lt_;
-  cudnnHandle_t cudnn_;
+  cublasLtHandle_t cublaslt_handle_{nullptr};
+  cudnnHandle_t cudnn_handle_{nullptr};
  std::unordered_map<int, CommandEncoder> encoders_;
 };

-Device& device(mlx::core::Device device);
+Device& device(int cuda_device);
+Device& device(mlx::core::Device d);
 CommandEncoder& get_command_encoder(Stream s);

-// Return an execution policy that does not sync for result.
-// Note that not all thrust APIs support async policy, confirm before using.
-inline auto thrust_policy(cudaStream_t stream) {
-  // TODO: Connect thrust's custom allocator with mlx's allocator.
-  return thrust::cuda::par_nosync.on(stream);
-}
-
 } // namespace mlx::core::cu
@@ -19,7 +19,7 @@ struct FloorDivide {
    if constexpr (cuda::std::is_integral_v<T>) {
      return x / y;
    } else {
-      return truncf(x / y);
+      return cuda::std::trunc(x / y);
    }
  }
 };
@@ -47,7 +47,7 @@ struct Remainder {
    } else if constexpr (is_complex_v<T>) {
      return x % y;
    } else {
-      T r = fmod(x, y);
+      T r = cuda::std::fmod(x, y);
      if (r != 0 && (r < 0 != y < 0)) {
        r = r + y;
      }
@@ -66,6 +66,7 @@ struct Equal {
 struct NaNEqual {
  template <typename T>
  __device__ bool operator()(T x, T y) {
+    using cuda::std::isnan;
    if constexpr (is_complex_v<T>) {
      return x == y ||
          (isnan(x.real()) && isnan(y.real()) && isnan(x.imag()) &&
@@ -110,8 +111,8 @@ struct LogAddExp {
  template <typename T>
  __device__ T operator()(T x, T y) {
    if constexpr (is_complex_v<T>) {
-      if (isnan(x.real()) || isnan(x.imag()) || isnan(y.real()) ||
-          isnan(y.imag())) {
+      if (cuda::std::isnan(x.real()) || cuda::std::isnan(x.imag()) ||
+          cuda::std::isnan(y.real()) || cuda::std::isnan(y.imag())) {
        return {
            cuda::std::numeric_limits<float>::quiet_NaN(),
            cuda::std::numeric_limits<float>::quiet_NaN()};
@@ -120,7 +121,7 @@ struct LogAddExp {
      auto min = x.real() < y.real() ? x : y;
      auto min_real = min.real();
      auto max_real = max.real();
-      if (!isfinite(min_real) && (min_real == max_real)) {
+      if (!cuda::std::isfinite(min_real) && (min_real == max_real)) {
        if (min_real < 0) {
          return min;
        } else {
@@ -130,7 +131,7 @@ struct LogAddExp {
        return Log1p{}(Exp{}(min - max)) + max;
      }
    } else {
-      if (isnan(x) || isnan(y)) {
+      if (cuda::std::isnan(x) || cuda::std::isnan(y)) {
        return cuda::std::numeric_limits<T>::quiet_NaN();
      }
      T maxval = max(x, y);
@@ -138,7 +139,7 @@ struct LogAddExp {
      return (minval == -cuda::std::numeric_limits<T>::infinity() ||
              maxval == cuda::std::numeric_limits<T>::infinity())
          ? maxval
-          : T(float(maxval) + log1p(expf(minval - maxval)));
+          : T(maxval + cuda::std::log1p(cuda::std::exp(minval - maxval)));
    }
  };
 };
@@ -149,12 +150,12 @@ struct Maximum {
    if constexpr (cuda::std::is_integral_v<T>) {
      return max(x, y);
    } else if constexpr (is_complex_v<T>) {
-      if (isnan(x.real()) || isnan(x.imag())) {
+      if (cuda::std::isnan(x.real()) || cuda::std::isnan(x.imag())) {
        return x;
      }
      return x > y ? x : y;
    } else {
-      if (isnan(x)) {
+      if (cuda::std::isnan(x)) {
        return x;
      }
      return x > y ? x : y;
@@ -168,12 +169,12 @@ struct Minimum {
    if constexpr (cuda::std::is_integral_v<T>) {
      return min(x, y);
    } else if constexpr (is_complex_v<T>) {
-      if (isnan(x.real()) || isnan(x.imag())) {
+      if (cuda::std::isnan(x.real()) || cuda::std::isnan(x.imag())) {
        return x;
      }
      return x < y ? x : y;
    } else {
-      if (isnan(x)) {
+      if (cuda::std::isnan(x)) {
        return x;
      }
      return x < y ? x : y;
@@ -219,9 +220,9 @@ struct Power {
      }
      return res;
    } else if constexpr (is_complex_v<T>) {
-      return pow(base, exp);
+      return cuda::std::pow(base, exp);
    } else {
-      return powf(base, exp);
+      return cuda::std::pow(base, exp);
    }
  }
 };
@@ -285,7 +286,7 @@ struct RightShift {
 struct ArcTan2 {
  template <typename T>
  __device__ T operator()(T y, T x) {
-    return atan2f(y, x);
+    return cuda::std::atan2(y, x);
  }
 };

@@ -4,98 +4,14 @@

 #include <cuda_bf16.h>
 #include <cuda_fp16.h>
-#include <cuda/std/limits>
 #include <cuda/std/type_traits>

 namespace mlx::core::cu {

-///////////////////////////////////////////////////////////////////////////////
-// Unary ops for half types.
-///////////////////////////////////////////////////////////////////////////////
-
-#if CUDART_VERSION < 12000 && __CUDA_ARCH__ < 800
-#define MLX_DEFINE_UNARY_OP(NAME, HALF_OP)           \
-  template <typename T>                              \
-  __forceinline__ __device__ auto NAME(T x) {        \
-    if constexpr (cuda::std::is_same_v<T, __half>) { \
-      return HALF_OP(x);                             \
-    } else {                                         \
-      return ::NAME(x);                              \
-    }                                                \
-  }
-#else
-#define MLX_DEFINE_UNARY_OP(NAME, HALF_OP)                         \
-  template <typename T>                                            \
-  __forceinline__ __device__ auto NAME(T x) {                      \
-    if constexpr (cuda::std::is_same_v<T, __half>) {               \
-      return HALF_OP(x);                                           \
-    } else if constexpr (cuda::std::is_same_v<T, __nv_bfloat16>) { \
-      return HALF_OP(x);                                           \
-    } else {                                                       \
-      return ::NAME(x);                                            \
-    }                                                              \
-  }
-#endif
-
-#define MLX_DEFINE_UNARY_OP_FALLBCK(NAME)                          \
-  template <typename T>                                            \
-  __forceinline__ __device__ auto NAME(T x) {                      \
-    if constexpr (cuda::std::is_same_v<T, __half>) {               \
-      return ::NAME(__half2float(x));                              \
-    } else if constexpr (cuda::std::is_same_v<T, __nv_bfloat16>) { \
-      return ::NAME(__bfloat162float(x));                          \
-    } else {                                                       \
-      return ::NAME(x);                                            \
-    }                                                              \
-  }
-
-MLX_DEFINE_UNARY_OP(abs, __habs)
-MLX_DEFINE_UNARY_OP(ceil, hceil)
-MLX_DEFINE_UNARY_OP(cos, hcos)
-MLX_DEFINE_UNARY_OP(exp, hexp)
-MLX_DEFINE_UNARY_OP(floor, hfloor)
-MLX_DEFINE_UNARY_OP(isnan, __hisnan)
-MLX_DEFINE_UNARY_OP(log, hlog)
-MLX_DEFINE_UNARY_OP(log2, hlog2)
-MLX_DEFINE_UNARY_OP(log10, hlog10)
-MLX_DEFINE_UNARY_OP(rint, hrint)
-MLX_DEFINE_UNARY_OP(rsqrt, hrsqrt)
-MLX_DEFINE_UNARY_OP(sin, hsin)
-MLX_DEFINE_UNARY_OP(sqrt, hsqrt)
-MLX_DEFINE_UNARY_OP_FALLBCK(acos)
-MLX_DEFINE_UNARY_OP_FALLBCK(acosh)
-MLX_DEFINE_UNARY_OP_FALLBCK(asin)
-MLX_DEFINE_UNARY_OP_FALLBCK(asinh)
-MLX_DEFINE_UNARY_OP_FALLBCK(atan)
-MLX_DEFINE_UNARY_OP_FALLBCK(atanh)
-MLX_DEFINE_UNARY_OP_FALLBCK(cosh)
-MLX_DEFINE_UNARY_OP_FALLBCK(log1p)
-MLX_DEFINE_UNARY_OP_FALLBCK(sinh)
-MLX_DEFINE_UNARY_OP_FALLBCK(tan)
-#if __CUDA_ARCH__ >= 1280
-MLX_DEFINE_UNARY_OP(tanh, htanh)
-#else
-MLX_DEFINE_UNARY_OP_FALLBCK(tanh)
-#endif
-
-#undef MLX_DEFINE_UNARY_OP
-#undef MLX_DEFINE_UNARY_OP_FALLBCK
-
 ///////////////////////////////////////////////////////////////////////////////
 // Binary ops for half types.
 ///////////////////////////////////////////////////////////////////////////////

-#if CUDART_VERSION < 12000 && __CUDA_ARCH__ < 800
-#define MLX_DEFINE_BINARY_OP(NAME, HALF_OP)          \
-  template <typename T>                              \
-  __forceinline__ __device__ auto NAME(T x, T y) {   \
-    if constexpr (cuda::std::is_same_v<T, __half>) { \
-      return HALF_OP(x, y);                          \
-    } else {                                         \
-      return ::NAME(x, y);                           \
-    }                                                \
-  }
-#else
 #define MLX_DEFINE_BINARY_OP(NAME, HALF_OP)                        \
  template <typename T>                                            \
  __forceinline__ __device__ auto NAME(T x, T y) {                 \
@@ -107,26 +23,12 @@ MLX_DEFINE_UNARY_OP_FALLBCK(tanh)
      return ::NAME(x, y);                                         \
    }                                                              \
  }
-#endif

 MLX_DEFINE_BINARY_OP(max, __hmax)
 MLX_DEFINE_BINARY_OP(min, __hmin)

 #undef MLX_DEFINE_BINARY_OP

-template <typename T>
-__forceinline__ __device__ T fmod(T x, T y) {
-  if constexpr (cuda::std::is_same_v<T, __half>) {
-    return __float2half(::fmod(__half2float(x), __half2float(y)));
-#if CUDART_VERSION >= 12000 || __CUDA_ARCH__ >= 800
-  } else if constexpr (cuda::std::is_same_v<T, __nv_bfloat16>) {
-    return __float2bfloat16(::fmod(__bfloat162float(x), __bfloat162float(y)));
-#endif
-  } else {
-    return ::fmod(x, y);
-  }
-}
-
 ///////////////////////////////////////////////////////////////////////////////
 // Additional C++ operator overrides between half types and native types.
 ///////////////////////////////////////////////////////////////////////////////
@@ -0,0 +1,184 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device/utils.cuh"
+
+namespace mlx::core::cu {
+
+__device__ __forceinline__ void hadamard_radix_m(float* x);
+
+template <int N>
+struct Pow2Log2 {
+  static_assert(
+      (N > 0) && ((N & (N - 1)) == 0),
+      "N must be a positive power of two.");
+  static constexpr int value = 1 + Pow2Log2<N / 2>::value;
+};
+
+template <>
+struct Pow2Log2<1> {
+  static constexpr int value = 0;
+};
+
+template <int R>
+__device__ __forceinline__ void hadamard_radix_pow2(float* x) {
+  constexpr int kLogR = Pow2Log2<R>::value;
+  int h = 1;
+#pragma unroll
+  for (int s = 0; s < kLogR; ++s) {
+#pragma unroll
+    for (int i = 0; i < R / 2; ++i) {
+      int k = i & (h - 1);
+      int j = ((i - k) << 1) + k;
+      float a = x[j];
+      float b = x[j + h];
+      x[j] = a + b;
+      x[j + h] = a - b;
+    }
+    h <<= 1;
+  }
+}
+
+template <typename T, int N, int max_radix, int read_width, int stride = 1>
+__global__ void
+hadamard_n(const T* in, T* out, float scale, long long num_transforms) {
+  constexpr int kNumThreads = N / max_radix;
+  constexpr int kLogN = Pow2Log2<N>::value;
+  constexpr int kLogR = Pow2Log2<max_radix>::value;
+  constexpr int kNumSteps = kLogN / kLogR;
+  constexpr int kLogFinal = kLogN % kLogR;
+  constexpr int kFinalRadix = 1 << kLogFinal;
+
+  if (threadIdx.x >= kNumThreads) {
+    return;
+  }
+
+  __shared__ T buf[N];
+  int i = threadIdx.x;
+
+  for (long long transform = blockIdx.x; transform < num_transforms;
+       transform += gridDim.x) {
+    long long base = (transform / stride) * static_cast<long long>(N) * stride +
+        (transform % stride);
+
+    if constexpr (stride == 1) {
+#pragma unroll
+      for (int j = 0; j < max_radix / read_width; ++j) {
+        int index = j * read_width * kNumThreads + i * read_width;
+#pragma unroll
+        for (int r = 0; r < read_width; ++r) {
+          buf[index + r] = in[base + index + r];
+        }
+      }
+    } else {
+#pragma unroll
+      for (int j = 0; j < max_radix; ++j) {
+        buf[j * kNumThreads + i] = in[base + (j * kNumThreads + i) * stride];
+      }
+    }
+    __syncthreads();
+
+    float x[max_radix];
+    int h = 1;
+
+#pragma unroll
+    for (int s = 0; s < kNumSteps; ++s) {
+      int k = i & (h - 1);
+      int j = ((i - k) << kLogR) + k;
+
+#pragma unroll
+      for (int r = 0; r < max_radix; ++r) {
+        x[r] = static_cast<float>(buf[j + h * r]);
+      }
+
+      hadamard_radix_pow2<max_radix>(x);
+
+#pragma unroll
+      for (int r = 0; r < max_radix; ++r) {
+        buf[j + h * r] = static_cast<T>(x[r]);
+      }
+
+      h <<= kLogR;
+      __syncthreads();
+    }
+
+    if constexpr (kFinalRadix > 1) {
+#pragma unroll
+      for (int t = 0; t < max_radix / kFinalRadix; ++t) {
+        int index = i + t * kNumThreads;
+        int k = index & (h - 1);
+        int j = ((index - k) << kLogFinal) + k;
+#pragma unroll
+        for (int r = 0; r < kFinalRadix; ++r) {
+          x[r] = static_cast<float>(buf[j + h * r]);
+        }
+
+        hadamard_radix_pow2<kFinalRadix>(x);
+
+#pragma unroll
+        for (int r = 0; r < kFinalRadix; ++r) {
+          buf[j + h * r] = static_cast<T>(x[r]);
+        }
+      }
+      __syncthreads();
+    }
+
+    if constexpr (stride == 1) {
+#pragma unroll
+      for (int j = 0; j < max_radix / read_width; ++j) {
+        int index = j * read_width * kNumThreads + i * read_width;
+#pragma unroll
+        for (int r = 0; r < read_width; ++r) {
+          float val = static_cast<float>(buf[index + r]);
+          out[base + index + r] = static_cast<T>(val * scale);
+        }
+      }
+    } else {
+#pragma unroll
+      for (int j = 0; j < max_radix; ++j) {
+        out[base + (j * kNumThreads + i) * stride] = buf[j * kNumThreads + i];
+      }
+    }
+
+    __syncthreads();
+  }
+}
+
+template <typename T, int N, int M, int read_width>
+__global__ void
+hadamard_m(const T* in, T* out, float scale, long long num_tasks) {
+  constexpr int kTasksPerBatch = N / read_width;
+
+  for (long long task = blockIdx.x * blockDim.x + threadIdx.x; task < num_tasks;
+       task += blockDim.x * gridDim.x) {
+    long long i = task % kTasksPerBatch;
+    long long batch = task / kTasksPerBatch;
+    long long base = batch * static_cast<long long>(M) * N;
+
+    float x[read_width][M];
+#pragma unroll
+    for (int c = 0; c < M; ++c) {
+#pragma unroll
+      for (int r = 0; r < read_width; ++r) {
+        x[r][c] = static_cast<float>(in[base + c * N + i * read_width + r]);
+      }
+    }
+
+#pragma unroll
+    for (int r = 0; r < read_width; ++r) {
+      hadamard_radix_m(x[r]);
+    }
+
+#pragma unroll
+    for (int c = 0; c < M; ++c) {
+#pragma unroll
+      for (int r = 0; r < read_width; ++r) {
+        out[base + c * N + i * read_width + r] =
+            static_cast<T>(x[r][c] * scale);
+      }
+    }
+  }
+}
+
+} // namespace mlx::core::cu
@@ -2,12 +2,12 @@

 #pragma once

-#include <cuda_fp8.h>
-
 #include "mlx/backend/cuda/device/fp16_math.cuh"
 #include "mlx/backend/cuda/device/utils.cuh"

+#include <cuda_fp8.h>
 #include <math_constants.h>
+#include <cuda/std/cmath>

 namespace mlx::core::cu {

@@ -17,7 +17,7 @@ struct Abs {
    if constexpr (cuda::std::is_unsigned_v<T>) {
      return x;
    } else {
-      return abs(x);
+      return cuda::std::abs(x);
    }
  }
 };
@@ -25,42 +25,42 @@ struct Abs {
 struct ArcCos {
  template <typename T>
  __device__ T operator()(T x) {
-    return acos(x);
+    return cuda::std::acos(x);
  }
 };

 struct ArcCosh {
  template <typename T>
  __device__ T operator()(T x) {
-    return acosh(x);
+    return cuda::std::acosh(x);
  }
 };

 struct ArcSin {
  template <typename T>
  __device__ T operator()(T x) {
-    return asin(x);
+    return cuda::std::asin(x);
  }
 };

 struct ArcSinh {
  template <typename T>
  __device__ T operator()(T x) {
-    return asinh(x);
+    return cuda::std::asinh(x);
  }
 };

 struct ArcTan {
  template <typename T>
  __device__ T operator()(T x) {
-    return atan(x);
+    return cuda::std::atan(x);
  }
 };

 struct ArcTanh {
  template <typename T>
  __device__ T operator()(T x) {
-    return atanh(x);
+    return cuda::std::atanh(x);
  }
 };

@@ -77,9 +77,9 @@ struct Ceil {
    if constexpr (cuda::std::is_integral_v<T>) {
      return x;
    } else if constexpr (is_complex_v<T>) {
-      return T{ceil(x.real()), ceil(x.imag())};
+      return T{cuda::std::ceil(x.real()), cuda::std::ceil(x.imag())};
    } else {
-      return ceil(x);
+      return cuda::std::ceil(x);
    }
  }
 };
@@ -87,21 +87,21 @@ struct Ceil {
 struct Conjugate {
  template <typename T>
  __device__ complex_t<T> operator()(complex_t<T> x) {
-    return conj(x);
+    return cuda::std::conj(x);
  }
 };

 struct Cos {
  template <typename T>
  __device__ T operator()(T x) {
-    return cos(x);
+    return cuda::std::cos(x);
  }
 };

 struct Cosh {
  template <typename T>
  __device__ T operator()(T x) {
-    return cosh(x);
+    return cuda::std::cosh(x);
  }
 };

@@ -134,20 +134,14 @@ struct ErfInv {
 struct Exp {
  template <typename T>
  __device__ T operator()(T x) {
-    return exp(x);
+    return cuda::std::exp(x);
  }
 };

 struct Expm1 {
  template <typename T>
  __device__ T operator()(T x) {
-    if constexpr (cuda::std::is_same_v<T, __half>) {
-      return expm1(__half2float(x));
-    } else if constexpr (cuda::std::is_same_v<T, __nv_bfloat16>) {
-      return expm1(__bfloat162float(x));
-    } else {
-      return expm1(x);
-    }
+    return cuda::std::expm1(x);
  }
 };

@@ -157,9 +151,9 @@ struct Floor {
    if constexpr (cuda::std::is_integral_v<T>) {
      return x;
    } else if constexpr (is_complex_v<T>) {
-      return T{floor(x.real()), floor(x.imag())};
+      return T{cuda::std::floor(x.real()), cuda::std::floor(x.imag())};
    } else {
-      return floor(x);
+      return cuda::std::floor(x);
    }
  }
 };
@@ -174,7 +168,7 @@ struct Imag {
 struct Log {
  template <typename T>
  __device__ T operator()(T x) {
-    return log(x);
+    return cuda::std::log(x);
  }
 };

@@ -185,7 +179,7 @@ struct Log2 {
      auto y = Log{}(x);
      return {y.real() / CUDART_LN2_F, y.imag() / CUDART_LN2_F};
    } else {
-      return log2(x);
+      return cuda::std::log2(x);
    }
  }
 };
@@ -193,7 +187,7 @@ struct Log2 {
 struct Log10 {
  template <typename T>
  __device__ T operator()(T x) {
-    return log10(x);
+    return cuda::std::log10(x);
  }
 };

@@ -216,7 +210,7 @@ struct Log1p {
        return {logf(z0), theta};
      }
    } else {
-      return log1p(z);
+      return cuda::std::log1p(z);
    }
  }
 };
@@ -249,9 +243,9 @@ struct Round {
  template <typename T>
  __device__ T operator()(T x) {
    if constexpr (is_complex_v<T>) {
-      return {rint(x.real()), rint(x.imag())};
+      return {cuda::std::rint(x.real()), cuda::std::rint(x.imag())};
    } else {
-      return rint(x);
+      return cuda::std::rint(x);
    }
  }
 };
@@ -259,7 +253,7 @@ struct Round {
 struct Sigmoid {
  template <typename T>
  __device__ T operator()(T x) {
-    T y = 1 / (1 + exp(abs(x)));
+    T y = 1 / (1 + cuda::std::exp(cuda::std::abs(x)));
    return (x < 0) ? y : 1 - y;
  }
 };
@@ -286,14 +280,14 @@ struct Sign {
 struct Sin {
  template <typename T>
  __device__ T operator()(T x) {
-    return sin(x);
+    return cuda::std::sin(x);
  }
 };

 struct Sinh {
  template <typename T>
  __device__ T operator()(T x) {
-    return sinh(x);
+    return cuda::std::sinh(x);
  }
 };

@@ -307,7 +301,7 @@ struct Square {
 struct Sqrt {
  template <typename T>
  __device__ T operator()(T x) {
-    return sqrt(x);
+    return cuda::std::sqrt(x);
  }
 };

@@ -316,6 +310,10 @@ struct Rsqrt {
  __device__ T operator()(T x) {
    if constexpr (is_complex_v<T>) {
      return 1.0f / Sqrt{}(x);
+    } else if constexpr (cuda::std::is_same_v<T, __half>) {
+      return rsqrt(__half2float(x));
+    } else if constexpr (cuda::std::is_same_v<T, __nv_bfloat16>) {
+      return rsqrt(__bfloat162float(x));
    } else {
      return rsqrt(x);
    }
@@ -325,14 +323,14 @@ struct Rsqrt {
 struct Tan {
  template <typename T>
  __device__ T operator()(T x) {
-    return tan(x);
+    return cuda::std::tan(x);
  }
 };

 struct Tanh {
  template <typename T>
  __device__ T operator()(T x) {
-    return tanh(x);
+    return cuda::std::tanh(x);
  }
 };

@@ -0,0 +1,232 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/gpu/device_info.h"
+#include "mlx/backend/cuda/cuda.h"
+
+#include <cuda_runtime.h>
+#include <dlfcn.h>
+
+#include <string>
+#include <unordered_map>
+#include <variant>
+#include <vector>
+
+namespace mlx::core {
+
+namespace {
+
+// NVML dynamic loading for accurate memory reporting
+// (cudaMemGetInfo only sees current process)
+
+typedef int nvmlReturn_t;
+typedef struct nvmlDevice_st* nvmlDevice_t;
+struct nvmlMemory_t {
+  unsigned long long total;
+  unsigned long long free;
+  unsigned long long used;
+};
+
+struct NVMLState {
+  void* handle = nullptr;
+  nvmlReturn_t (*nvmlInit_v2)() = nullptr;
+  nvmlReturn_t (*nvmlDeviceGetHandleByUUID)(const char*, nvmlDevice_t*) =
+      nullptr;
+  nvmlReturn_t (*nvmlDeviceGetMemoryInfo)(nvmlDevice_t, nvmlMemory_t*) =
+      nullptr;
+};
+
+bool nvml_init(NVMLState& nvml) {
+#ifdef _WIN32
+  nvml.handle = dlopen("nvml.dll", RTLD_LAZY);
+  if (!nvml.handle) {
+    nvml.handle = dlopen(
+        "C:\\Program Files\\NVIDIA Corporation\\NVSMI\\nvml.dll", RTLD_LAZY);
+  }
+#else
+  nvml.handle = dlopen("libnvidia-ml.so.1", RTLD_LAZY);
+#endif
+  if (!nvml.handle)
+    return false;
+
+  nvml.nvmlInit_v2 =
+      (decltype(nvml.nvmlInit_v2))dlsym(nvml.handle, "nvmlInit_v2");
+  nvml.nvmlDeviceGetHandleByUUID =
+      (decltype(nvml.nvmlDeviceGetHandleByUUID))dlsym(
+          nvml.handle, "nvmlDeviceGetHandleByUUID");
+  nvml.nvmlDeviceGetMemoryInfo = (decltype(nvml.nvmlDeviceGetMemoryInfo))dlsym(
+      nvml.handle, "nvmlDeviceGetMemoryInfo");
+
+  if (!nvml.nvmlInit_v2 || !nvml.nvmlDeviceGetHandleByUUID ||
+      !nvml.nvmlDeviceGetMemoryInfo) {
+    return false;
+  }
+  return nvml.nvmlInit_v2() == 0;
+}
+
+bool nvml_get_memory(
+    NVMLState& nvml,
+    const char* uuid,
+    size_t* free,
+    size_t* total) {
+  if (!nvml.handle)
+    return false;
+  nvmlDevice_t device;
+  if (nvml.nvmlDeviceGetHandleByUUID(uuid, &device) != 0)
+    return false;
+  nvmlMemory_t mem;
+  if (nvml.nvmlDeviceGetMemoryInfo(device, &mem) != 0)
+    return false;
+  *free = mem.free;
+  *total = mem.total;
+  return true;
+}
+
+std::string format_uuid(const cudaUUID_t& uuid) {
+  char buf[64];
+  snprintf(
+      buf,
+      sizeof(buf),
+      "GPU-%02x%02x%02x%02x-%02x%02x-%02x%02x-%02x%02x-%02x%02x%02x%02x%02x%02x",
+      (unsigned char)uuid.bytes[0],
+      (unsigned char)uuid.bytes[1],
+      (unsigned char)uuid.bytes[2],
+      (unsigned char)uuid.bytes[3],
+      (unsigned char)uuid.bytes[4],
+      (unsigned char)uuid.bytes[5],
+      (unsigned char)uuid.bytes[6],
+      (unsigned char)uuid.bytes[7],
+      (unsigned char)uuid.bytes[8],
+      (unsigned char)uuid.bytes[9],
+      (unsigned char)uuid.bytes[10],
+      (unsigned char)uuid.bytes[11],
+      (unsigned char)uuid.bytes[12],
+      (unsigned char)uuid.bytes[13],
+      (unsigned char)uuid.bytes[14],
+      (unsigned char)uuid.bytes[15]);
+  return buf;
+}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info_impl(int device_index) {
+  // Static cache of device properties including UUID (needed for NVML lookup)
+  static auto all_devices = []() {
+    // Get device count
+    int count = 0;
+    cudaGetDeviceCount(&count);
+
+    // Collect info for all devices
+    struct DeviceInfo {
+      std::unordered_map<std::string, std::variant<std::string, size_t>> info;
+      std::string uuid;
+    };
+
+    std::vector<DeviceInfo> devices;
+
+    for (int i = 0; i < count; ++i) {
+      cudaDeviceProp prop;
+      cudaGetDeviceProperties(&prop, i);
+
+      DeviceInfo dev;
+      dev.info["device_name"] = std::string(prop.name);
+      dev.uuid = format_uuid(prop.uuid);
+      dev.info["uuid"] = dev.uuid;
+
+      // Architecture string (e.g., "sm_89")
+      char arch[16];
+      snprintf(arch, sizeof(arch), "sm_%d%d", prop.major, prop.minor);
+      dev.info["architecture"] = std::string(arch);
+
+      // PCI bus ID (domain:bus:device.function)
+      char pci_id[32];
+      snprintf(
+          pci_id,
+          sizeof(pci_id),
+          "%04x:%02x:%02x.0",
+          prop.pciDomainID,
+          prop.pciBusID,
+          prop.pciDeviceID);
+      dev.info["pci_bus_id"] = std::string(pci_id);
+
+      // Compute capability as size_t (to match Metal's variant type)
+      dev.info["compute_capability_major"] = static_cast<size_t>(prop.major);
+      dev.info["compute_capability_minor"] = static_cast<size_t>(prop.minor);
+
+      devices.push_back(std::move(dev));
+    }
+    return devices;
+  }();
+
+  // Initialize NVML once for fresh memory reads
+  static NVMLState nvml;
+  static bool nvml_initialized = nvml_init(nvml);
+
+  if (device_index < 0 ||
+      device_index >= static_cast<int>(all_devices.size())) {
+    static auto empty =
+        std::unordered_map<std::string, std::variant<std::string, size_t>>();
+    return empty;
+  }
+
+  // Return a copy with fresh memory info
+  // Using thread_local to avoid locks while keeping free_memory fresh
+  thread_local auto device_info_copy =
+      std::unordered_map<std::string, std::variant<std::string, size_t>>();
+
+  device_info_copy = all_devices[device_index].info;
+
+  // Get fresh memory info - try NVML first (system-wide), fallback to
+  // cudaMemGetInfo (process-level)
+  size_t free_mem, total_mem;
+
+  if (nvml_initialized &&
+      nvml_get_memory(
+          nvml,
+          all_devices[device_index].uuid.c_str(),
+          &free_mem,
+          &total_mem)) {
+    // NVML succeeded - use system-wide memory
+  } else {
+    // Fallback to cudaMemGetInfo (process-scoped)
+    int prev_device;
+    cudaGetDevice(&prev_device);
+    cudaSetDevice(device_index);
+    cudaMemGetInfo(&free_mem, &total_mem);
+    cudaSetDevice(prev_device);
+  }
+
+  device_info_copy["free_memory"] = free_mem;
+  device_info_copy["total_memory"] = total_mem;
+
+  return device_info_copy;
+}
+
+} // anonymous namespace
+
+namespace gpu {
+
+bool is_available() {
+  return true;
+}
+
+int device_count() {
+  int count = 0;
+  cudaGetDeviceCount(&count);
+  return count;
+}
+
+const std::unordered_map<std::string, std::variant<std::string, size_t>>&
+device_info(int device_index) {
+  return device_info_impl(device_index);
+}
+
+} // namespace gpu
+
+namespace cu {
+
+bool is_available() {
+  return true;
+}
+
+} // namespace cu
+
+} // namespace mlx::core
@@ -3,7 +3,6 @@
 #include "mlx/backend/gpu/eval.h"
 #include "mlx/backend/cuda/allocator.h"
 #include "mlx/backend/cuda/device.h"
-#include "mlx/backend/gpu/available.h"
 #include "mlx/primitives.h"
 #include "mlx/scheduler.h"

@@ -11,10 +10,6 @@

 namespace mlx::core::gpu {

-bool is_available() {
-  return true;
-}
-
 void new_stream(Stream s) {
  // Force initalization of CUDA, so CUDA runtime get destroyed at last.
  cudaFree(nullptr);
@@ -26,6 +21,10 @@ void new_stream(Stream s) {

 void eval(array& arr) {
  nvtx3::scoped_range r("gpu::eval");
+  // Ensure CUDA context is active on this thread. Required when MLX is called
+  // from threads that have not yet established a CUDA context (e.g. thread
+  // pools, language runtimes that migrate work across OS threads).
+  cu::device(arr.primitive().stream().device).make_current();
  auto outputs = arr.outputs();
  {
    // If the array is a tracer hold a reference
@@ -3,6 +3,7 @@
 #include "mlx/backend/cuda/allocator.h"
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/event.h"
+#include "mlx/backend/gpu/device_info.h"
 #include "mlx/event.h"
 #include "mlx/scheduler.h"

@@ -119,9 +120,10 @@ void CudaEvent::init_pool() {
 class CopyableCudaEvent {
 public:
  explicit CopyableCudaEvent(Device& d)
-      : event_(std::make_shared<CudaEvent>(
-            d,
-            cudaEventDisableTiming | cudaEventBlockingSync)) {}
+      : event_(
+            std::make_shared<CudaEvent>(
+                d,
+                cudaEventDisableTiming | cudaEventBlockingSync)) {}

  void wait() {
    event_->wait();
@@ -172,45 +174,95 @@ class CopyableCudaEvent {
 // AtomicEvent implementations
 ///////////////////////////////////////////////////////////////////////////////

-__host__ __device__ void event_wait(AtomicEvent::Atomic* ac, uint64_t value) {
-  uint64_t current;
-  while ((current = ac->load()) < value) {
-    ac->wait(current);
+__host__ __device__ void event_wait(uint32_t* ptr, uint32_t value) {
+  cuda::atomic_ref<uint32_t> ac(*ptr);
+  uint32_t current;
+  while ((current = ac.load()) < value) {
+    ac.wait(current);
  }
 }

-__host__ __device__ void event_signal(AtomicEvent::Atomic* ac, uint64_t value) {
-  ac->store(value);
-  ac->notify_all();
+__host__ __device__ void event_signal(uint32_t* ptr, uint32_t value) {
+  cuda::atomic_ref<uint32_t> ac(*ptr);
+  ac.store(value);
+  ac.notify_all();
 }

-__global__ void event_wait_kernel(AtomicEvent::Atomic* ac, uint64_t value) {
-  event_wait(ac, value);
+__global__ void event_wait_kernel(uint32_t* ptr, uint32_t value) {
+  event_wait(ptr, value);
 }

-__global__ void event_signal_kernel(AtomicEvent::Atomic* ac, uint64_t value) {
-  event_signal(ac, value);
+__global__ void event_signal_kernel(uint32_t* ptr, uint32_t value) {
+  __threadfence_system();
+  event_signal(ptr, value);
+  __threadfence_system();
 }

-AtomicEvent::AtomicEvent() {
-  buf_ = std::shared_ptr<Buffer>(
-      new Buffer{allocator().malloc(sizeof(Atomic))}, [](Buffer* ptr) {
-        allocator().free(*ptr);
-        delete ptr;
-      });
-  *static_cast<uint64_t*>(buf_->raw_ptr()) = 0;
+auto check_gpu_coherency() {
+  static auto coherency = []() {
+    int device_count = gpu::device_count();
+    bool concurrent_managed_access = true;
+    bool host_native_atomic = true;
+    for (int i = 0; i < device_count; ++i) {
+      auto& d = cu::device(i);
+      concurrent_managed_access &= d.concurrent_managed_access();
+      host_native_atomic &= d.host_native_atomic();
+    }
+    return std::make_tuple(concurrent_managed_access, host_native_atomic);
+  }();
+  return coherency;
 }

-void AtomicEvent::wait(uint64_t value) {
+AtomicEvent::AtomicEvent(Device& d) {
+  void* buf;
+  cudaError_t (*cuda_free)(void*);
+  // There are 3 kinds of systems we are implementing for:
+  // 1. concurrentManagedAccess == true
+  //    => use cuda::atom_ref on managed memory
+  // 2. hostNativeAtomicSupported == true
+  //    => use cuda::atom_ref on pinned host memory
+  // 2. no hardware cpu/gpu coherency
+  //    => use cuda::atom_ref on device memory
+  d.make_current();
+  auto [concurrent_managed_access, host_native_atomic] = check_gpu_coherency();
+  if (concurrent_managed_access) {
+    CHECK_CUDA_ERROR(cudaMallocManaged(&buf, sizeof(uint32_t)));
+    cuda_free = cudaFree;
+    coherent_ = true;
+  } else if (host_native_atomic) {
+    CHECK_CUDA_ERROR(cudaMallocHost(&buf, sizeof(uint32_t)));
+    cuda_free = cudaFreeHost;
+    coherent_ = true;
+  } else {
+    CHECK_CUDA_ERROR(cudaMalloc(&buf, sizeof(uint32_t)));
+    cuda_free = cudaFree;
+    coherent_ = false;
+  }
+  buf_ = std::shared_ptr<void>(
+      buf, [cuda_free](void* buf) { CHECK_CUDA_ERROR(cuda_free(buf)); });
+  if (coherent_) {
+    *ptr() = 0;
+  } else {
+    CHECK_CUDA_ERROR(cudaMemset(buf, 0, sizeof(uint32_t)));
+  }
+}
+
+void AtomicEvent::wait(uint32_t value) {
  nvtx3::scoped_range r("cu::AtomicEvent::wait");
-  event_wait(atomic(), value);
+  if (coherent_) {
+    event_wait(ptr(), value);
+  } else {
+    while (!is_signaled(value)) {
+      std::this_thread::yield();
+    }
+  }
 }

-void AtomicEvent::wait(cudaStream_t stream, uint64_t value) {
-  event_wait_kernel<<<1, 1, 0, stream>>>(atomic(), value);
+void AtomicEvent::wait(cudaStream_t stream, uint32_t value) {
+  event_wait_kernel<<<1, 1, 0, stream>>>(ptr(), value);
 }

-void AtomicEvent::wait(Stream s, uint64_t value) {
+void AtomicEvent::wait(Stream s, uint32_t value) {
  nvtx3::scoped_range r("cu::AtomicEvent::wait(s)");
  if (s.device == mlx::core::Device::cpu) {
    scheduler::enqueue(s, [*this, value]() mutable { wait(value); });
@@ -222,22 +274,26 @@ void AtomicEvent::wait(Stream s, uint64_t value) {
  }
 }

-void AtomicEvent::signal(uint64_t value) {
+void AtomicEvent::signal(uint32_t value) {
  nvtx3::scoped_range r("cu::AtomicEvent::signal");
-  event_signal(atomic(), value);
+  if (coherent_) {
+    event_signal(ptr(), value);
+  } else {
+    signal(signal_stream(), value);
+  }
 }

-void AtomicEvent::signal(cudaStream_t stream, uint64_t value) {
-  event_signal_kernel<<<1, 1, 0, stream>>>(atomic(), value);
+void AtomicEvent::signal(cudaStream_t stream, uint32_t value) {
+  event_signal_kernel<<<1, 1, 0, stream>>>(ptr(), value);
 }

-void AtomicEvent::signal(Stream s, uint64_t value) {
+void AtomicEvent::signal(Stream s, uint32_t value) {
  nvtx3::scoped_range r("cu::AtomicEvent::signal(s)");
  if (s.device == mlx::core::Device::cpu) {
    // Signal through a GPU stream so the atomic is updated in GPU - updating
    // the atomic in CPU sometimes does not get GPU notified.
-    static CudaStream stream(device(mlx::core::Device::gpu));
-    scheduler::enqueue(s, [*this, value]() mutable { signal(stream, value); });
+    scheduler::enqueue(
+        s, [*this, value]() mutable { signal(signal_stream(), value); });
  } else {
    auto& encoder = get_command_encoder(s);
    encoder.commit();
@@ -246,14 +302,26 @@ void AtomicEvent::signal(Stream s, uint64_t value) {
  }
 }

-bool AtomicEvent::is_signaled(uint64_t value) const {
-  nvtx3::scoped_range r("cu::AtomicEvent::is_signaled");
-  return atomic()->load() >= value;
+bool AtomicEvent::is_signaled(uint32_t val) const {
+  return value() >= val;
 }

-uint64_t AtomicEvent::value() const {
+uint32_t AtomicEvent::value() const {
  nvtx3::scoped_range r("cu::AtomicEvent::value");
-  return atomic()->load();
+  if (coherent_) {
+    cuda::atomic_ref<uint32_t> ac(*ptr());
+    return ac.load();
+  } else {
+    uint32_t val;
+    CHECK_CUDA_ERROR(
+        cudaMemcpy(&val, ptr(), sizeof(uint32_t), cudaMemcpyDeviceToHost));
+    return val;
+  }
+}
+
+const CudaStream& AtomicEvent::signal_stream() {
+  static CudaStream stream(device(0));
+  return stream;
 }

 } // namespace cu
@@ -280,11 +348,12 @@ struct EventImpl {
    if (is_created()) {
      return;
    }
+    auto& d = cu::device(s.device);
    if (s.device == mlx::core::Device::cpu || signal_value > 1) {
      nvtx3::mark("Using slow AtomicEvent");
-      atomic = std::make_unique<cu::AtomicEvent>();
+      atomic = std::make_unique<cu::AtomicEvent>(d);
    } else {
-      cuda = std::make_unique<cu::CopyableCudaEvent>(cu::device(s.device));
+      cuda = std::make_unique<cu::CopyableCudaEvent>(d);
    }
  }
 };
@@ -54,25 +54,26 @@ class CudaEvent {
 // CudaEvent so the latter should always be preferred when possible.
 class AtomicEvent {
 public:
-  using Atomic = cuda::atomic<uint64_t>;
+  AtomicEvent(Device& d);

-  AtomicEvent();
-
-  void wait(uint64_t value);
-  void wait(cudaStream_t stream, uint64_t value);
-  void wait(Stream s, uint64_t value);
-  void signal(uint64_t value);
-  void signal(cudaStream_t stream, uint64_t value);
-  void signal(Stream s, uint64_t value);
-  bool is_signaled(uint64_t value) const;
-  uint64_t value() const;
+  void wait(uint32_t value);
+  void wait(cudaStream_t stream, uint32_t value);
+  void wait(Stream s, uint32_t value);
+  void signal(uint32_t value);
+  void signal(cudaStream_t stream, uint32_t value);
+  void signal(Stream s, uint32_t value);
+  bool is_signaled(uint32_t value) const;
+  uint32_t value() const;

 private:
-  Atomic* atomic() const {
-    return static_cast<AtomicEvent::Atomic*>(buf_->raw_ptr());
+  const CudaStream& signal_stream();
+
+  uint32_t* ptr() const {
+    return static_cast<uint32_t*>(buf_.get());
  }

-  std::shared_ptr<allocator::Buffer> buf_;
+  bool coherent_;
+  std::shared_ptr<void> buf_;
 };

 } // namespace mlx::core::cu
@@ -14,7 +14,8 @@ struct FenceImpl {

 Fence::Fence(Stream s) {
  fence_ = std::shared_ptr<void>(
-      new FenceImpl{0}, [](void* ptr) { delete static_cast<FenceImpl*>(ptr); });
+      new FenceImpl{0, cu::device(s.device)},
+      [](void* ptr) { delete static_cast<FenceImpl*>(ptr); });
 }

 void Fence::wait(Stream s, const array&) {
@@ -29,15 +30,9 @@ void Fence::update(Stream s, const array& a, bool cross_device) {
    auto& cbuf =
        *static_cast<cu::CudaBuffer*>(const_cast<array&>(a).buffer().ptr());
    if (cbuf.device != -1) {
-      void* new_data;
-      CHECK_CUDA_ERROR(cudaMallocManaged(&new_data, cbuf.size));
-      cbuf.device = -1;
-      auto& encoder = cu::device(s.device).get_command_encoder(s);
+      auto& encoder = cu::get_command_encoder(s);
      encoder.commit();
-      CHECK_CUDA_ERROR(cudaMemcpyAsync(
-          new_data, cbuf.data, cbuf.size, cudaMemcpyDefault, encoder.stream()));
-      CHECK_CUDA_ERROR(cudaFreeAsync(cbuf.data, encoder.stream()));
-      cbuf.data = new_data;
+      cu::allocator().move_to_unified_memory(cbuf, encoder.stream());
    }
  }
  fence->count++;
@@ -27,8 +27,9 @@ cublasComputeType_t dtype_to_compute_type(Dtype dtype) {
      return mlx::core::env::enable_tf32() ? CUBLAS_COMPUTE_32F_FAST_TF32
                                           : CUBLAS_COMPUTE_32F;
    default:
-      throw std::runtime_error(fmt::format(
-          "Unsupported dtype in CublasGemm: {}.", dtype_to_string(dtype)));
+      throw std::runtime_error(
+          fmt::format(
+              "Unsupported dtype in CublasGemm: {}.", dtype_to_string(dtype)));
  }
 }

@@ -72,6 +73,14 @@ CublasGemm::CublasGemm(
      batch_count,
      a_batch_stride,
      b_batch_stride);
+
+  // alpha and beta are both host pointers
+  cublasLtPointerMode_t pointer_mode = CUBLASLT_POINTER_MODE_HOST;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_POINTER_MODE,
+      &pointer_mode,
+      sizeof(pointer_mode)));
 }

 CublasGemm::CublasGemm(
@@ -214,8 +223,8 @@ void CublasGemm::execute(
    const void* a,
    const void* b,
    const void* c,
-    float alpha /* = 1 */,
-    float beta /* = 0 */) {
+    const float alpha /* = 1 */,
+    const float beta /* = 0 */) {
  const void* alpha_ptr = &alpha;
  const void* beta_ptr = &beta;
  complex64_t alpha_c, beta_c;
@@ -182,7 +182,6 @@ void CublasGemm::run_batched(
          cu::set_mm_device_pointers_nd<ndim_constant()>,
          num_blocks,
          block_dims,
-          0,
          gpu_ptr<int8_t*>(pointers),
          gpu_ptr<int8_t>(a),
          gpu_ptr<int8_t>(b),
@@ -199,7 +198,6 @@ void CublasGemm::run_batched(
        cu::set_mm_device_pointers_g,
        num_blocks,
        block_dims,
-        0,
        gpu_ptr<int8_t*>(pointers),
        gpu_ptr<int8_t>(a),
        gpu_ptr<int8_t>(b),
@@ -270,7 +268,6 @@ void CublasGemm::run_batched(
          cu::set_addmm_device_pointers_nd<ndim_constant()>,
          num_blocks,
          block_dims,
-          0,
          gpu_ptr<int8_t*>(pointers),
          gpu_ptr<int8_t>(a),
          gpu_ptr<int8_t>(b),
@@ -289,7 +286,6 @@ void CublasGemm::run_batched(
        cu::set_addmm_device_pointers_g,
        num_blocks,
        block_dims,
-        0,
        gpu_ptr<int8_t*>(pointers),
        gpu_ptr<int8_t>(a),
        gpu_ptr<int8_t>(b),
@@ -104,6 +104,68 @@ __global__ void gemv_batched(
      mat + mat_offset, vec + vec_offset, out + batch_idx * rows, rows, cols);
 }

+template <typename T, int rows_per_block, int n_per_thread>
+__global__ void gemv_gather(
+    const T* mat,
+    const T* vec,
+    T* out,
+    uint32_t* mat_indices,
+    uint32_t* vec_indices,
+    int rows,
+    int cols,
+    const __grid_constant__ Shape mat_batch_shape,
+    const __grid_constant__ Strides mat_batch_strides,
+    int mat_batch_ndim,
+    const __grid_constant__ Shape vec_batch_shape,
+    const __grid_constant__ Strides vec_batch_strides,
+    int vec_batch_ndim,
+    const __grid_constant__ Shape index_shape,
+    const __grid_constant__ Strides mat_index_strides,
+    const __grid_constant__ Strides vec_index_strides,
+    int index_batch_ndim) {
+  auto block = cg::this_thread_block();
+  auto indices_idx = block.group_index().y;
+  uint32_t index_mat, index_vec;
+  if (index_batch_ndim > 1) {
+    auto [mat_idx_offset, vec_idx_offset] = elem_to_loc(
+        indices_idx,
+        index_shape.data(),
+        mat_index_strides.data(),
+        vec_index_strides.data(),
+        index_batch_ndim);
+    index_mat = mat_indices[mat_idx_offset];
+    index_vec = vec_indices[vec_idx_offset];
+  } else {
+    index_mat = mat_indices[indices_idx * mat_index_strides[0]];
+    index_vec = vec_indices[indices_idx * vec_index_strides[0]];
+  }
+
+  int64_t mat_offset;
+  if (mat_batch_ndim > 1) {
+    mat_offset = elem_to_loc(
+        index_mat,
+        mat_batch_shape.data(),
+        mat_batch_strides.data(),
+        mat_batch_ndim);
+  } else {
+    mat_offset = index_mat * mat_batch_strides[0];
+  }
+
+  int64_t vec_offset;
+  if (vec_batch_ndim > 1) {
+    vec_offset = elem_to_loc(
+        index_vec,
+        vec_batch_shape.data(),
+        vec_batch_strides.data(),
+        vec_batch_ndim);
+  } else {
+    vec_offset = index_vec * vec_batch_strides[0];
+  }
+
+  gemv_impl<T, rows_per_block, n_per_thread>(
+      mat + mat_offset, vec + vec_offset, out + indices_idx * rows, rows, cols);
+}
+
 bool can_use_gemv(int M, int N, int K, bool a_transposed, bool b_transposed) {
  return K % 32 == 0 && ((M == 1 && b_transposed) || (N == 1 && !a_transposed));
 }
@@ -174,7 +236,6 @@ void gemv(
            kernel,
            num_blocks_x,
            block_dims,
-            0,
            mat,
            vec,
            gpu_ptr<DataType>(out),
@@ -186,7 +247,6 @@ void gemv(
            kernel,
            dim3{num_blocks_x, batch_count},
            block_dims,
-            0,
            mat,
            vec,
            gpu_ptr<DataType>(out),
@@ -201,4 +261,64 @@ void gemv(
  });
 }

+void gather_mv(
+    const array& mat_,
+    const array& vec_,
+    const array& mat_indices,
+    const array& vec_indices,
+    array& out,
+    int N,
+    int K,
+    CommandEncoder& encoder) {
+  encoder.set_input_array(mat_);
+  encoder.set_input_array(vec_);
+  encoder.set_input_array(mat_indices);
+  encoder.set_input_array(vec_indices);
+  encoder.set_output_array(out);
+  dispatch_inexact_types(out.dtype(), "gather_mv", [&](auto type_tag) {
+    using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    dim3 block_dims{WARP_SIZE, rows_per_block};
+    int rows = N;
+    int cols = K;
+    uint32_t batch_size = static_cast<uint32_t>(out.size() / N);
+    const DataType* mat = gpu_ptr<DataType>(mat_);
+    const DataType* vec = gpu_ptr<DataType>(vec_);
+
+    uint32_t num_blocks_x = (rows + rows_per_block - 1) / rows_per_block;
+    int n_per_t;
+    if (K % 128 == 0 && is_aligned<4>(mat) && is_aligned<4>(vec)) {
+      n_per_t = 4;
+    } else if (K % 64 == 0 && is_aligned<2>(mat) && is_aligned<2>(vec)) {
+      n_per_t = 2;
+    } else {
+      n_per_t = 1;
+    }
+
+    dispatch_n_per_thread(n_per_t, [&](auto n_per_thread) {
+      auto kernel = gemv_gather<DataType, rows_per_block, n_per_thread()>;
+      encoder.add_kernel_node(
+          kernel,
+          dim3{num_blocks_x, batch_size},
+          block_dims,
+          mat,
+          vec,
+          gpu_ptr<DataType>(out),
+          gpu_ptr<uint32_t>(mat_indices),
+          gpu_ptr<uint32_t>(vec_indices),
+          rows,
+          cols,
+          const_param(mat_.shape()),
+          const_param(mat_.strides()),
+          mat_.ndim() - 2,
+          const_param(vec_.shape()),
+          const_param(vec_.strides()),
+          vec_.ndim() - 2,
+          const_param(mat_indices.shape()),
+          const_param(mat_indices.strides()),
+          const_param(vec_indices.strides()),
+          mat_indices.ndim());
+    });
+  });
+}
+
 } // namespace mlx::core::cu
@@ -21,4 +21,14 @@ void gemv(
    const mlx::core::Strides& b_batch_strides,
    CommandEncoder& encoder);

+void gather_mv(
+    const array& mat,
+    const array& vec,
+    const array& mat_indices,
+    const array& vec_indices,
+    array& out,
+    int N,
+    int K,
+    CommandEncoder& encoder);
+
 } // namespace mlx::core::cu
@@ -22,4 +22,18 @@ void cutlass_grouped_gemm_unaligned(
    array& out,
    cu::CommandEncoder& encoder);

+void cutlass_segmented_mm(
+    bool a_transposed,
+    int lda,
+    bool b_transposed,
+    int ldb,
+    int num_segments,
+    int M,
+    int N,
+    const array& a,
+    const array& b,
+    const array& segments,
+    array& out,
+    cu::CommandEncoder& encoder);
+
 } // namespace mlx::core
@@ -1,5 +1,7 @@
 // Copyright © 2025 Apple Inc.

+#include "mlx/backend/cuda/cublas_utils.h"
+#include "mlx/backend/cuda/cutlass_utils.cuh"
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/gemms/grouped_gemm.h"
 #include "mlx/backend/cuda/kernel_utils.cuh"
@@ -9,7 +11,6 @@
 #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 {
@@ -92,11 +93,127 @@ __global__ void prepare_grouped_mm_data(
  }
 }

+__global__ void prepare_segmented_mm_data(
+    const uint32_t* segments,
+    int num_segments,
+    int M,
+    int N,
+    int lda,
+    int ldb,
+    int item_size,
+    bool a_transposed,
+    bool b_transposed,
+    int8_t* a_start,
+    int8_t* b_start,
+    int8_t* out_start,
+    ProblemSize* problem_sizes,
+    int64_t* a_lds,
+    int64_t* b_lds,
+    int64_t* out_lds,
+    void** a_ptrs,
+    void** b_ptrs,
+    void** out_ptrs) {
+  int idx = cg::this_grid().thread_rank();
+  if (idx >= num_segments)
+    return;
+
+  int64_t start = segments[2 * idx];
+  int64_t end = segments[2 * idx + 1];
+  int K_i = (end > start) ? static_cast<int>(end - start) : 0;
+
+  problem_sizes[idx] = {M, N, K_i};
+  a_lds[idx] = lda;
+  b_lds[idx] = ldb;
+  out_lds[idx] = N;
+
+  // Offset into K dimension depends on layout:
+  // A [M,K]: row-major offset = start, col-major offset = start * lda
+  // B [K,N]: row-major offset = start * ldb, col-major offset = start
+  int64_t a_offset = a_transposed ? start * lda : start;
+  int64_t b_offset = b_transposed ? start : start * ldb;
+
+  a_ptrs[idx] = a_start + a_offset * item_size;
+  b_ptrs[idx] = b_start + b_offset * item_size;
+  out_ptrs[idx] = out_start + static_cast<int64_t>(idx) * M * N * item_size;
+}
+
 } // namespace cu

 namespace {

-template <typename T, int kAlignment, typename Arch, typename OpClass>
+// Shared GEMM configuration for every type and arch.
+template <typename T, typename ArchTag, int kAlignmentC>
+struct CommonGemmConfiguration {
+  using Element = T;
+  using Arch = ArchTag;
+  using Accumulator = std::conditional_t<(sizeof(T) < 4), float, T>;
+  using EpilogueOutputOp = cutlass::epilogue::thread::
+      LinearCombination<T, kAlignmentC, Accumulator, Accumulator>;
+};
+
+// Slow GEMM configuration as fallback.
+template <
+    typename T,
+    typename Arch,
+    int kAlignmentC = 1,
+    bool kEnableTF32 = false,
+    typename Enable = void>
+struct GemmConfiguration : public CommonGemmConfiguration<T, Arch, 1> {
+  using OpClass = cutlass::arch::OpClassSimt;
+  using ThreadblockShape = cutlass::gemm::GemmShape<128, 128, 8>;
+  using WarpShape = cutlass::gemm::GemmShape<32, 64, 8>;
+  using InstructionShape = cutlass::gemm::GemmShape<1, 1, 1>;
+  static const int kAlignmentAB = 1;
+  static const int kStages = 2;
+};
+
+// Specialized GEMM configuration for sm80 and later.
+template <typename T, typename Arch, int kAlignmentC>
+struct GemmConfiguration<
+    T,
+    Arch,
+    kAlignmentC,
+    true,
+    std::enable_if_t<Arch::kMinComputeCapability >= 80 && sizeof(T) <= 4>>
+    : public CommonGemmConfiguration<T, cutlass::arch::Sm80, kAlignmentC> {
+  using OpClass = cutlass::arch::OpClassTensorOp;
+  using ThreadblockShape = cutlass::gemm::GemmShape<256, 128, 32>;
+  using WarpShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using InstructionShape = cutlass::gemm::GemmShape<16, 8, 32 / sizeof(T)>;
+  static const int kAlignmentAB = 1;
+  static const int kStages = 2;
+};
+
+// Specialized GEMM configuration for tf32 on sm80.
+template <int kAlignmentC>
+struct GemmConfiguration<float, cutlass::arch::Sm80, kAlignmentC, true>
+    : public CommonGemmConfiguration<float, cutlass::arch::Sm80, kAlignmentC> {
+  using OpClass = cutlass::arch::OpClassTensorOp;
+  using ThreadblockShape = cutlass::gemm::GemmShape<256, 128, 32>;
+  using WarpShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using InstructionShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static const int kAlignmentAB = 1;
+  static const int kStages = 3; // use SM80_CP_ASYNC
+};
+
+// Get direct access to kernel.
+template <typename GemmKernel>
+class GemmGroupedEncoder
+    : public cutlass::gemm::device::GemmGrouped<GemmKernel> {
+ public:
+  void encode(cu::CommandEncoder& encoder) {
+    encoder.add_kernel_node_ex(
+        cutlass::Kernel<GemmKernel>,
+        {static_cast<uint32_t>(this->params_.threadblock_count), 1, 1},
+        {GemmKernel::kThreadCount, 1, 1},
+        {},
+        sizeof(typename GemmKernel::SharedStorage),
+        this->params_);
+  }
+};
+
+// Invoke the grouped GEMM of CUTLASS 2.x API, which supports small alignments.
+template <typename GemmConfiguration>
 void grouped_gemm_v2(
    bool a_transposed,
    bool b_transposed,
@@ -109,79 +226,92 @@ void grouped_gemm_v2(
    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,
+          a_transposed_tag.value,
          cutlass::layout::ColumnMajor,
          cutlass::layout::RowMajor>;
      using LayoutB = std::conditional_t<
-          b_transposed_tag,
+          b_transposed_tag.value,
          cutlass::layout::ColumnMajor,
          cutlass::layout::RowMajor>;
      using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmGrouped<
-          T,
+          typename GemmConfiguration::Element,
          LayoutA,
          cutlass::ComplexTransform::kNone,
-          kAlignment,
-          T,
+          GemmConfiguration::kAlignmentAB,
+          typename GemmConfiguration::Element,
          LayoutB,
          cutlass::ComplexTransform::kNone,
-          kAlignment,
-          T,
+          GemmConfiguration::kAlignmentAB,
+          typename GemmConfiguration::Element,
          cutlass::layout::RowMajor,
-          ElementAccumulator,
-          OpClass,
-          Arch,
+          typename GemmConfiguration::Accumulator,
+          typename GemmConfiguration::OpClass,
+          typename GemmConfiguration::Arch,
          typename GemmConfiguration::ThreadblockShape,
          typename GemmConfiguration::WarpShape,
          typename GemmConfiguration::InstructionShape,
-          EpilogueOutputOp,
+          typename GemmConfiguration::EpilogueOutputOp,
          cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,
          GemmConfiguration::kStages>::GemmKernel;
-      using GemmGrouped =
-          typename cutlass::gemm::device::GemmGrouped<GemmKernel>;
+      using GemmGrouped = GemmGroupedEncoder<GemmKernel>;

-      typename EpilogueOutputOp::Params epilogue_op(
-          /* alpha */ 1, /* beta */ 0);
+      static int threadblock_count = GemmGrouped::sufficient();
      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),
+          threadblock_count,
+          {/* alpha */ 1, /* beta */ 0},
+          reinterpret_cast<typename GemmGrouped::ElementA**>(a_ptrs),
+          reinterpret_cast<typename GemmGrouped::ElementB**>(b_ptrs),
+          reinterpret_cast<typename GemmGrouped::ElementC**>(out_ptrs),
+          reinterpret_cast<typename GemmGrouped::ElementC**>(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)));
-      }
+      CHECK_CUTLASS_ERROR(gemm.initialize(
+          args,
+          allocate_workspace(encoder, gemm.get_workspace_size(args)),
+          encoder.stream()));
+      gemm.encode(encoder);
    });
  });
 }

+template <typename F>
+void dispatch_cutlass_arch(cu::Device& device, F&& f) {
+  if (device.compute_capability_major() < 8) {
+    f(type_identity<cutlass::arch::Sm75>{});
+  } else if (device.compute_capability_major() == 8) {
+    f(type_identity<cutlass::arch::Sm80>{});
+  } else {
+    f(type_identity<cutlass::arch::Sm90>{});
+  }
+}
+
+auto* get_grouped_mm_funcion(Dtype dtype, int N, cu::Device& device) {
+  auto* fun = grouped_gemm_v2<GemmConfiguration<float, cutlass::arch::Sm75>>;
+  dispatch_float_types(dtype, "grouped_gemm_v2", [&](auto type_tag) {
+    using DataType = cutlass_type_t<MLX_GET_TYPE(type_tag)>;
+    dispatch_cutlass_arch(device, [&](auto arch_tag) {
+      using Arch = MLX_GET_TYPE(arch_tag);
+      dispatch_bool(N % 8 == 0, [&](auto is_out_aligned) {
+        constexpr int kAlignmentC = is_out_aligned ? 8 : 1;
+        dispatch_bool(env::enable_tf32(), [&](auto kEnableTF32) {
+          fun = grouped_gemm_v2<
+              GemmConfiguration<DataType, Arch, kAlignmentC, kEnableTF32>>;
+        });
+      });
+    });
+  });
+  return fun;
+}
+
 } // namespace

 void cutlass_grouped_gemm_unaligned(
@@ -195,6 +325,9 @@ void cutlass_grouped_gemm_unaligned(
    const array& indices,
    array& out,
    cu::CommandEncoder& encoder) {
+  int K = a.shape(-1);
+  int N = b.shape(-1);
+
  // Prepare device pointers for matmul.
  int problem_sizes_nbytes =
      group_count * cuda::ceil_div(sizeof(ProblemSize), 8) * 8;
@@ -214,23 +347,24 @@ void cutlass_grouped_gemm_unaligned(

  // Fill the pointers by computing offsets from indices.
  constexpr int N_READS = 4;
-  size_t n_threads = cuda::ceil_div(indices.size(), N_READS);
+  int 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 block_dims(std::min(n_threads, 1024));
  dim3 num_blocks(1);

  encoder.set_input_array(indices);
  encoder.set_output_array(gemm_args);
-  encoder.add_kernel_node(
+  encoder.add_kernel_node_ex(
      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,
+      K,
+      N,
      lda,
      ldb,
      out.itemsize(),
@@ -249,29 +383,11 @@ void cutlass_grouped_gemm_unaligned(
      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);
+  auto* fun = get_grouped_mm_funcion(a.dtype(), N, encoder.device());
  fun(a_transposed,
      b_transposed,
      group_count,
@@ -285,4 +401,85 @@ void cutlass_grouped_gemm_unaligned(
      encoder);
 }

+void cutlass_segmented_mm(
+    bool a_transposed,
+    int lda,
+    bool b_transposed,
+    int ldb,
+    int num_segments,
+    int M,
+    int N,
+    const array& a,
+    const array& b,
+    const array& segments,
+    array& out,
+    cu::CommandEncoder& encoder) {
+  // Allocate grouped GEMM args on device.
+  int problem_sizes_nbytes =
+      num_segments * cuda::ceil_div(sizeof(ProblemSize), 8) * 8;
+  int nbytes = problem_sizes_nbytes +
+      num_segments * (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 + num_segments;
+  int64_t* out_lds = b_lds + num_segments;
+  void** a_ptrs = reinterpret_cast<void**>(out_lds + num_segments);
+  void** b_ptrs = a_ptrs + num_segments;
+  void** out_ptrs = b_ptrs + num_segments;
+
+  // Build problem descriptions from segments on the GPU.
+  int block_size = std::min(num_segments, 256);
+  int num_blocks = cuda::ceil_div(num_segments, block_size);
+
+  encoder.set_input_array(segments);
+  encoder.set_output_array(gemm_args);
+  encoder.add_kernel_node_ex(
+      cu::prepare_segmented_mm_data,
+      dim3(num_blocks),
+      dim3(block_size),
+      {},
+      0,
+      gpu_ptr<uint32_t>(segments),
+      num_segments,
+      M,
+      N,
+      static_cast<int>(lda),
+      static_cast<int>(ldb),
+      static_cast<int>(out.itemsize()),
+      a_transposed,
+      b_transposed,
+      gpu_ptr<int8_t>(a),
+      gpu_ptr<int8_t>(b),
+      gpu_ptr<int8_t>(out),
+      problem_sizes,
+      a_lds,
+      b_lds,
+      out_lds,
+      a_ptrs,
+      b_ptrs,
+      out_ptrs);
+
+  // Dispatch grouped GEMM.
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(gemm_args);
+  encoder.set_output_array(out);
+  auto* fun = get_grouped_mm_funcion(a.dtype(), N, encoder.device());
+  fun(a_transposed,
+      b_transposed,
+      num_segments,
+      problem_sizes,
+      a_lds,
+      b_lds,
+      out_lds,
+      a_ptrs,
+      b_ptrs,
+      out_ptrs,
+      encoder);
+}
+
 } // namespace mlx::core
@@ -0,0 +1,245 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/common/hadamard.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/jit_module.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/dtype_utils.h"
+#include "mlx/primitives.h"
+
+#include <fmt/format.h>
+#include <nvtx3/nvtx3.hpp>
+
+#include <algorithm>
+#include <cassert>
+#include <sstream>
+#include <stdexcept>
+#include <string_view>
+
+namespace mlx::core {
+
+namespace {
+
+constexpr int MAX_HADAMARD_THREADS_PER_BLOCK = 256;
+
+std::string gen_hadamard_codelet(int m) {
+  std::ostringstream source;
+  source << "namespace mlx::core::cu {\n";
+  source << "__device__ __forceinline__ void hadamard_radix_m(float* x) {\n";
+  if (m == 1) {
+    source << "}\n";
+    source << "} // namespace mlx::core::cu\n";
+    return source.str();
+  }
+
+  auto h_matrices = hadamard_matrices();
+  auto it = h_matrices.find(m);
+  if (it == h_matrices.end()) {
+    throw std::runtime_error("[hadamard] Invalid radix m.");
+  }
+  auto& matrix = it->second;
+
+  source << "  float tmp[" << m << "];\n";
+  auto start = 1;
+  auto end = matrix.find('\n', start);
+  int row_idx = 0;
+  while (end != std::string_view::npos) {
+    auto row = matrix.substr(start, end - start);
+    source << "  tmp[" << row_idx << "] =";
+    for (int i = 0; i < row.length(); ++i) {
+      source << " " << row[i] << " x[" << i << "]";
+    }
+    source << ";\n";
+    start = end + 1;
+    end = matrix.find('\n', start);
+    row_idx++;
+  }
+  source << "  #pragma unroll\n";
+  source << "  for (int i = 0; i < " << m << "; ++i) { x[i] = tmp[i]; }\n";
+  source << "}\n";
+  source << "} // namespace mlx::core::cu\n";
+  return source.str();
+}
+
+std::string hadamard_n_kernel_name(
+    const Dtype& dtype,
+    int n,
+    int max_radix,
+    int read_width,
+    int stride) {
+  return fmt::format(
+      "mlx::core::cu::hadamard_n<{}, {}, {}, {}, {}>",
+      dtype_to_cuda_type(dtype),
+      n,
+      max_radix,
+      read_width,
+      stride);
+}
+
+std::string
+hadamard_m_kernel_name(const Dtype& dtype, int n, int m, int read_width) {
+  return fmt::format(
+      "mlx::core::cu::hadamard_m<{}, {}, {}, {}>",
+      dtype_to_cuda_type(dtype),
+      n,
+      m,
+      read_width);
+}
+
+void hadamard_mn_contiguous(
+    const array& x,
+    array& y,
+    int m,
+    int n1,
+    int n2,
+    float scale,
+    const Stream& s) {
+  const int n = n1 * n2;
+  const int read_width_n1 = (n1 == 2) ? 2 : 4;
+  const int read_width_n2 = (n2 == 2) ? 2 : 4;
+  const int read_width_m = (n == 2 || m == 28) ? 2 : 4;
+  const int max_radix_1 = std::min(n1, 16);
+  const int max_radix_2 = std::min(n2, 16);
+  const float scale_n1 = 1.0f;
+  const float scale_n2 = (m == 1) ? scale : 1.0f;
+  const float scale_m = scale;
+
+  const std::string n1_kernel_name =
+      hadamard_n_kernel_name(x.dtype(), n1, max_radix_1, read_width_n1, n2);
+  const std::string n2_kernel_name =
+      hadamard_n_kernel_name(x.dtype(), n2, max_radix_2, read_width_n2, 1);
+  const std::string m_kernel_name =
+      hadamard_m_kernel_name(x.dtype(), n, m, read_width_m);
+
+  const std::string module_name = fmt::format(
+      "hadamard_{}_{}_{}_{}_{}_{}_{}_{}",
+      dtype_to_string(x.dtype()),
+      n,
+      m,
+      n1,
+      n2,
+      read_width_n1,
+      read_width_n2,
+      read_width_m);
+
+  cu::JitModule& mod = cu::get_jit_module(s.device, module_name, [&]() {
+    std::vector<std::string> kernel_names = {n2_kernel_name};
+    if (n1 > 1) {
+      kernel_names.push_back(n1_kernel_name);
+    }
+    if (m > 1) {
+      kernel_names.push_back(m_kernel_name);
+    }
+
+    std::string source = R"(
+        #include "mlx/backend/cuda/device/utils.cuh"
+    )";
+    source += gen_hadamard_codelet(m);
+    source += R"(
+        #include "mlx/backend/cuda/device/hadamard.cuh"
+    )";
+
+    return std::make_tuple(false, std::move(source), std::move(kernel_names));
+  });
+
+  auto& encoder = cu::get_command_encoder(s);
+
+  if (n1 > 1) {
+    const int64_t num_transforms = x.size() / n1;
+    const uint32_t num_blocks =
+        static_cast<uint32_t>(std::min<int64_t>(num_transforms, 65535));
+
+    encoder.set_input_array(x);
+    encoder.set_output_array(y);
+
+    cu::KernelArgs args;
+    args.append(x);
+    args.append(y);
+    args.append(scale_n1);
+    args.append(num_transforms);
+
+    auto kernel = mod.get_kernel(n1_kernel_name);
+    encoder.add_kernel_node_raw(
+        kernel, num_blocks, n1 / max_radix_1, {}, 0, args.args());
+  }
+
+  {
+    const auto& in = (n1 > 1) ? y : x;
+    const int64_t num_transforms = x.size() / n2;
+    const uint32_t num_blocks =
+        static_cast<uint32_t>(std::min<int64_t>(num_transforms, 65535));
+
+    encoder.set_input_array(in);
+    encoder.set_output_array(y);
+
+    cu::KernelArgs args;
+    args.append(in);
+    args.append(y);
+    args.append(scale_n2);
+    args.append(num_transforms);
+
+    auto kernel = mod.get_kernel(n2_kernel_name);
+    encoder.add_kernel_node_raw(
+        kernel, num_blocks, n2 / max_radix_2, {}, 0, args.args());
+  }
+
+  if (m > 1) {
+    const int64_t num_tasks = x.size() / (m * read_width_m);
+    const uint32_t block_dim = static_cast<uint32_t>(
+        std::min<int64_t>(num_tasks, MAX_HADAMARD_THREADS_PER_BLOCK));
+    const uint32_t num_blocks = static_cast<uint32_t>(
+        std::min<int64_t>((num_tasks + block_dim - 1) / block_dim, 65535));
+
+    encoder.set_input_array(y);
+    encoder.set_output_array(y);
+
+    cu::KernelArgs args;
+    args.append(y);
+    args.append(y);
+    args.append(scale_m);
+    args.append(num_tasks);
+
+    auto kernel = mod.get_kernel(m_kernel_name);
+    encoder.add_kernel_node_raw(
+        kernel, num_blocks, block_dim, {}, 0, args.args());
+  }
+}
+
+} // namespace
+
+void Hadamard::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("Hadamard::eval_gpu");
+  assert(inputs.size() == 1);
+
+  auto& in = inputs[0];
+  if (in.dtype() != float16 && in.dtype() != bfloat16 &&
+      in.dtype() != float32) {
+    throw std::invalid_argument("[hadamard] Unsupported type.");
+  }
+
+  // n = m * 2^k where m in (1, 12, 20, 28)
+  auto [n, m] = decompose_hadamard(in.shape().back());
+  int n1 = 1;
+  int n2 = n;
+  if (n > 8192) {
+    for (n2 = 2; n2 * n2 < n; n2 *= 2) {
+    }
+    n1 = n / n2;
+  }
+
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+  if (in.flags().row_contiguous) {
+    if (in.is_donatable()) {
+      out.copy_shared_buffer(in);
+    } else {
+      out.set_data(cu::malloc_async(out.nbytes(), encoder));
+    }
+    hadamard_mn_contiguous(in, out, m, n1, n2, scale_, s);
+  } else {
+    copy_gpu(in, out, CopyType::General, s);
+    hadamard_mn_contiguous(out, out, m, n1, n2, scale_, s);
+  }
+}
+
+} // namespace mlx::core
@@ -86,13 +86,14 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
    std::vector<std::string> kernel_names;
    for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) {
      for (int large = 0; large <= 1; ++large) {
-        kernel_names.push_back(fmt::format(
-            "mlx::core::cu::gather<{}, {}, {}, {}, {}>",
-            dtype_to_cuda_type(out.dtype()),
-            dtype_to_cuda_type(idx_dtype),
-            nidx,
-            ndim,
-            large ? "int64_t" : "int32_t"));
+        kernel_names.push_back(
+            fmt::format(
+                "mlx::core::cu::gather<{}, {}, {}, {}, {}>",
+                dtype_to_cuda_type(out.dtype()),
+                dtype_to_cuda_type(idx_dtype),
+                nidx,
+                ndim,
+                large ? "int64_t" : "int32_t"));
      }
    }
    return std::make_tuple(false, jit_source_gather, std::move(kernel_names));
@@ -129,7 +130,8 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {

  auto kernel = mod.get_kernel(kernel_name);
  auto [num_blocks, block_dims] = get_launch_args(out, large);
-  encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
+  encoder.add_kernel_node_raw(
+      kernel, num_blocks, block_dims, {}, 0, args.args());
 }

 void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
@@ -179,14 +181,15 @@ void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
    std::vector<std::string> kernel_names;
    for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) {
      for (int large = 0; large <= 1; ++large) {
-        kernel_names.push_back(fmt::format(
-            "mlx::core::cu::scatter<{}, {}, mlx::core::cu::Scatter{}, {}, {}, {}>",
-            dtype_to_cuda_type(out.dtype()),
-            dtype_to_cuda_type(idx_dtype),
-            op,
-            nidx,
-            ndim,
-            large ? "int64_t" : "int32_t"));
+        kernel_names.push_back(
+            fmt::format(
+                "mlx::core::cu::scatter<{}, {}, mlx::core::cu::Scatter{}, {}, {}, {}>",
+                dtype_to_cuda_type(out.dtype()),
+                dtype_to_cuda_type(idx_dtype),
+                op,
+                nidx,
+                ndim,
+                large ? "int64_t" : "int32_t"));
      }
    }
    return std::make_tuple(false, jit_source_scatter, std::move(kernel_names));
@@ -230,7 +233,8 @@ void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
  encoder.set_output_array(out);
  auto kernel = mod.get_kernel(kernel_name);
  auto [num_blocks, block_dims] = get_launch_args(upd, large);
-  encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
+  encoder.add_kernel_node_raw(
+      kernel, num_blocks, block_dims, {}, 0, args.args());
 }

 void GatherAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
@@ -258,14 +262,15 @@ void GatherAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
    for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) {
      for (int contiguous = 0; contiguous < 4; ++contiguous) {
        for (int large = 0; large <= 1; ++large) {
-          kernel_names.push_back(fmt::format(
-              "mlx::core::cu::gather_axis<{}, {}, {}, {}, {}, {}>",
-              dtype_to_cuda_type(out.dtype()),
-              dtype_to_cuda_type(idx.dtype()),
-              ndim,
-              contiguous & 1 ? true : false,
-              contiguous & 2 ? true : false,
-              large ? "int64_t" : "int32_t"));
+          kernel_names.push_back(
+              fmt::format(
+                  "mlx::core::cu::gather_axis<{}, {}, {}, {}, {}, {}>",
+                  dtype_to_cuda_type(out.dtype()),
+                  dtype_to_cuda_type(idx.dtype()),
+                  ndim,
+                  contiguous & 1 ? true : false,
+                  contiguous & 2 ? true : false,
+                  large ? "int64_t" : "int32_t"));
        }
      }
    }
@@ -319,7 +324,8 @@ void GatherAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
  encoder.set_output_array(out);
  auto kernel = mod.get_kernel(kernel_name);
  auto [num_blocks, block_dims] = get_launch_args(idx, large);
-  encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
+  encoder.add_kernel_node_raw(
+      kernel, num_blocks, block_dims, {}, 0, args.args());
 }

 void ScatterAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
@@ -360,15 +366,16 @@ void ScatterAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
    for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) {
      for (int contiguous = 0; contiguous < 4; ++contiguous) {
        for (int large = 0; large <= 1; ++large) {
-          kernel_names.push_back(fmt::format(
-              "mlx::core::cu::scatter_axis<{}, {}, mlx::core::cu::Scatter{}, {}, {}, {}, {}>",
-              dtype_to_cuda_type(out.dtype()),
-              dtype_to_cuda_type(idx.dtype()),
-              op,
-              ndim,
-              contiguous & 1 ? true : false,
-              contiguous & 2 ? true : false,
-              large ? "int64_t" : "int32_t"));
+          kernel_names.push_back(
+              fmt::format(
+                  "mlx::core::cu::scatter_axis<{}, {}, mlx::core::cu::Scatter{}, {}, {}, {}, {}>",
+                  dtype_to_cuda_type(out.dtype()),
+                  dtype_to_cuda_type(idx.dtype()),
+                  op,
+                  ndim,
+                  contiguous & 1 ? true : false,
+                  contiguous & 2 ? true : false,
+                  large ? "int64_t" : "int32_t"));
        }
      }
    }
@@ -424,7 +431,8 @@ void ScatterAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
  encoder.set_output_array(out);
  auto kernel = mod.get_kernel(kernel_name);
  auto [num_blocks, block_dims] = get_launch_args(idx, large);
-  encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
+  encoder.add_kernel_node_raw(
+      kernel, num_blocks, block_dims, {}, 0, args.args());
 }

 } // namespace mlx::core
@@ -12,7 +12,6 @@

 #include <fmt/format.h>
 #include <nvrtc.h>
-#include <unistd.h>

 namespace mlx::core::cu {

@@ -32,7 +31,10 @@ const std::vector<std::string>& include_path_args() {
  static std::vector<std::string> cached_args = []() {
    std::vector<std::string> args;
    // Add path to bundled CCCL headers.
-    auto root_dir = current_binary_dir().parent_path();
+    auto root_dir = current_binary_dir();
+#if !defined(_WIN32)
+    root_dir = root_dir.parent_path();
+#endif
    auto path = root_dir / "include" / "cccl";
 #if defined(MLX_CCCL_DIR)
    if (!std::filesystem::exists(path)) {
@@ -80,6 +82,18 @@ const std::filesystem::path& ptx_cache_dir() {
      cache =
          std::filesystem::temp_directory_path() / "mlx" / version() / "ptx";
    }
+
+#if defined(_WIN32)
+    // Add "\\?\" prefix to support long file path.
+    const wchar_t* long_path_prefix = L"\\\\?\\";
+    if (cache.is_relative()) {
+      cache = std::filesystem::absolute(cache);
+    }
+    if (!cache.native().starts_with(long_path_prefix)) {
+      cache = long_path_prefix + cache.native();
+    }
+#endif
+
    if (!std::filesystem::exists(cache)) {
      std::error_code error;
      if (!std::filesystem::create_directories(cache, error)) {
@@ -94,12 +108,7 @@ const std::filesystem::path& ptx_cache_dir() {
 std::filesystem::path get_ptx_path(
    const std::filesystem::path& cache_dir,
    const std::string& module_name) {
-#ifdef _WIN32
-  constexpr int max_file_name_length = 140;
-#else
  constexpr int max_file_name_length = 245;
-#endif
-
  if (module_name.size() <= max_file_name_length) {
    return cache_dir / (module_name + ".ptx");
  }
@@ -222,6 +231,7 @@ constexpr const char* g_include_names[] = {
    INCLUDE_PREFIX "config.h",
    INCLUDE_PREFIX "complex.cuh",
    INCLUDE_PREFIX "fp16_math.cuh",
+    INCLUDE_PREFIX "hadamard.cuh",
    INCLUDE_PREFIX "indexing.cuh",
    INCLUDE_PREFIX "scatter_ops.cuh",
    INCLUDE_PREFIX "unary_ops.cuh",
@@ -238,6 +248,7 @@ constexpr const char* g_headers[] = {
    jit_source_config,
    jit_source_complex,
    jit_source_fp16_math,
+    jit_source_hadamard,
    jit_source_indexing,
    jit_source_scatter_ops,
    jit_source_unary_ops,
@@ -272,7 +283,7 @@ void compile(
  std::vector<const char*> args;
  bool use_sass = compiler_supports_device_sass(device);
  auto cc = device.compute_capability_major();
-  std::string arch_tag = (cc == 90 || cc == 100 || cc == 121) ? "a" : "";
+  std::string arch_tag = (cc >= 9) ? "a" : "";
  std::string compute = fmt::format(
      "--gpu-architecture={}_{}{}{}",
      use_sass ? "sm" : "compute",
@@ -321,7 +332,7 @@ void load_module(
    const std::string& ptx,
    const std::vector<std::pair<std::string, std::string>>& ptx_kernels,
    CUmodule& module_,
-    std::unordered_map<std::string, std::tuple<CUfunction, bool, uint>>&
+    std::unordered_map<std::string, std::tuple<CUfunction, bool, uint32_t>>&
        kernels) {
  // Load module.
  char jit_log[4089] = {};
@@ -331,8 +342,9 @@ void load_module(
  CUresult jit_result = cuModuleLoadDataEx(
      &module_, ptx.data(), std::size(options), options, values);
  if (jit_result != CUDA_SUCCESS) {
-    throw std::runtime_error(fmt::format(
-        "Failed to load compiled {} kernel: {}.", module_name, jit_log));
+    throw std::runtime_error(
+        fmt::format(
+            "Failed to load compiled {} kernel: {}.", module_name, jit_log));
  }

  // Load kernels.
@@ -383,7 +395,7 @@ JitModule::~JitModule() {
  CHECK_CUDA_ERROR(cuModuleUnload(module_));
 }

-std::pair<CUfunction, uint> JitModule::get_kernel_and_dims(
+std::pair<CUfunction, uint32_t> JitModule::get_kernel_and_dims(
    const std::string& kernel_name,
    std::function<void(CUfunction)> configure_kernel) {
  auto it = kernels_.find(kernel_name);
@@ -99,13 +99,14 @@ class JitModule {
  CUfunction get_kernel(
      const std::string& kernel_name,
      std::function<void(CUfunction)> configure_kernel = nullptr);
-  std::pair<CUfunction, uint> get_kernel_and_dims(
+  std::pair<CUfunction, uint32_t> get_kernel_and_dims(
      const std::string& kernel_name,
      std::function<void(CUfunction)> configure_kernel = nullptr);

 private:
  CUmodule module_{nullptr};
-  std::unordered_map<std::string, std::tuple<CUfunction, bool, uint>> kernels_;
+  std::unordered_map<std::string, std::tuple<CUfunction, bool, uint32_t>>
+      kernels_;
 };

 std::unordered_map<std::string, JitModule>& get_jit_module_cache();
@@ -30,15 +30,15 @@ std::pair<dim3, dim3> get_grid_and_block(int dim0, int dim1, int dim2) {
  return std::make_pair(dim3(gx, gy, gz), dim3(bx, by, bz));
 }

-std::tuple<dim3, uint> get_launch_args(
+std::tuple<dim3, uint32_t> get_launch_args(
    size_t size,
    const Shape& shape,
    const Strides& strides,
    bool large,
    int work_per_thread /* = 1 */,
-    uint max_block_dim /* = 1024 */) {
+    uint32_t max_block_dim /* = 1024 */) {
  size_t nthreads = cuda::ceil_div(size, work_per_thread);
-  uint block_dim = max_block_dim < nthreads ? max_block_dim : nthreads;
+  uint32_t block_dim = max_block_dim < nthreads ? max_block_dim : nthreads;
  dim3 num_blocks;
  if (large) {
    num_blocks = get_2d_grid_dims(shape, strides, work_per_thread);
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Cheng	ce45c52505	[CUDA] Use qmv kernel for fp quantizations (#3239 )	2026-03-12 07:25:17 +09:00
Jagrit Digani	0879a6acba	Add initial tuning for M5 pro and max (#3211 )	2026-03-11 14:05:43 -07:00
Long Yixing	a9573f92f6	[CUDA] Implement SegmentedMM (#3238 )	2026-03-11 13:31:43 -07:00
Cheng	1c2d7041ab	Remove quantized_utils.cuh (#3237 )	2026-03-11 19:45:51 +09:00
Daniel Hiltgen	fd6d304b3a	win: fix cuda build (#3204 )	2026-03-11 12:58:04 +09:00
Anastasiia Filippova	e1e1399e1b	Hybrid sharding (#3194 )	2026-03-10 11:47:25 +01:00
Cheng	9d03a1b0d9	[CUDA] Support 3/5/6-bit quants in QMV (#3236 )	2026-03-10 19:09:48 +09:00
Cheng	8d022bcb86	Remove custom fp4/fp8 classes (#3212 )	2026-03-10 16:08:01 +09:00
Michelle DiMarco	d2702a4fc1	Fix non-strict module update with extra weights (#3214 ) Co-authored-by: Michelle DiMarco <m_dimarco@apple.com>	2026-03-09 22:03:50 -07:00
Dan Anderson	6ac5280db4	Fix assigning bool to float16/bfloat16 (#3229 ) Co-authored-by: KD2YCU <me@kd2ycu.com> Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2026-03-09 21:50:05 -07:00
Dan Anderson	572e0a4ac3	Validate dims in rope (#3230 ) Co-authored-by: KD2YCU <me@kd2ycu.com>	2026-03-09 21:48:45 -07:00
Dan Anderson	9bbd375eec	Fix return value in einsum_path for simple contractions (#3232 ) Co-authored-by: KD2YCU <me@kd2ycu.com> Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2026-03-09 21:48:26 -07:00
Dan Anderson	a25399cbd4	Validate num_splits in split (#3234 ) Co-authored-by: KD2YCU <me@kd2ycu.com>	2026-03-09 21:47:08 -07:00
Cheng	5a347b2ec8	[CUDA] Faster compilation and batch support in QMV (#3213 )	2026-03-10 13:45:10 +09:00
Dan Anderson	db487f3649	PR #3220 LayerNorm VJP returns zeros_like(weight) instead of zeros_like(bias placeholder) (#3231 ) Co-authored-by: KD2YCU <me@kd2ycu.com>	2026-03-09 17:06:12 -07:00
Dan Anderson	8f5ff2ea41	PR#3226 Fix (#3227 ) Co-authored-by: KD2YCU <me@kd2ycu.com>	2026-03-09 15:06:33 -07:00
Angelos Katharopoulos	d06c3c8936	Improve mlx.distributed_config (#3199 )	2026-03-09 13:17:51 -07:00
Long Yixing	be872ebdef	[CUDA] implement Hadamard transform (#3179 )	2026-03-05 09:34:19 +01:00
Cheng	3b3590bf5f	[CUDA] Use fp16 accumulation for 4-bit quant in GEMV (#3197 )	2026-03-05 07:58:23 +09:00
Cheng	3c565437a5	[CUDA] Quantized GEMV (#3180 )	2026-03-04 08:59:31 +09:00
dependabot[bot]	9eef9f1774	Bump actions/upload-artifact from 6 to 7 (#3188 ) Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>	2026-03-04 08:26:24 +09:00
dependabot[bot]	e320a2adc2	Bump actions/download-artifact from 7 to 8 (#3189 ) Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>	2026-03-04 08:24:12 +09:00
willem adnet	8cd377b7db	Add the bartlett function (#3155 )	2026-03-03 11:40:54 -08:00
Christophe Prat	f145ece976	Fix/missing libs in docs (#3190 )	2026-03-03 09:53:18 -08:00
AN Long	1ce7118303	Fix ref leak in mx.save/load with file like object (#3187 )	2026-03-02 16:55:29 -08:00
Anastasiia Filippova	72e04f7fb7	[CUDA] Fsdp (easy) (#3130 )	2026-03-01 23:29:09 +01:00
Cheng	6482d13dd3	Skip Hopper-only kernels in CI (#3184 )	2026-02-27 23:22:08 -08:00
Angelos Katharopoulos	1e3736b19d	Bump the patch version (#3185 )	2026-02-27 23:21:37 -08:00
Angelos Katharopoulos	365d6f29b4	Bump the minor version (#3183 )	2026-02-27 14:02:03 -08:00
Angelos Katharopoulos	d7a553c536	Enable passing in a GPU architecture string via env var (#3176 )	2026-02-27 11:37:53 -08:00
Robert Johansson	c8536f5248	Fix compile_fuse broadcast split aliasing bug (#3166 ) Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2026-02-26 18:17:28 -08:00
Cheng	0c8107ce8a	[CUDA] Heuristics for Hopper QMM (#3173 )	2026-02-27 09:22:10 +09:00
Angelos Katharopoulos	5c4abd2f06	JACCL refactor and small update (#3174 )	2026-02-26 13:56:19 -08:00
Anastasiia Filippova	4e00919e5c	[CUDA][NCCL] group split (#3172 )	2026-02-26 09:26:20 +01:00
Cheng	6ec0192270	Enable setting thread block cluster for Hopper and later (#3168 )	2026-02-26 16:15:16 +09:00
willem adnet	a8ba5ac3e0	Implement mlx.core.blackman (#3136 )	2026-02-25 13:42:40 -08:00
Cheng	6304c285d3	[CUDA] FPxINT quantized matmul for Hopper (#3160 )	2026-02-25 09:10:18 +09:00
Awni Hannun	cb198268d5	[Metal] Fix event leak (#3159 )	2026-02-23 19:50:48 -08:00
Gleb Sterkin	1d8d693d08	[Metal] Add implicit matmul pathway for mx.conv3d (#3147 ) Co-authored-by: Gleb Sterkin <g_sterkin@apple.com> Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2026-02-23 17:52:50 -08:00
Kellen Sun	d4c81062ad	[Metal] Fix 32-bit integer overflow in conv3d unfold kernel (#3143 )	2026-02-19 10:01:22 -08:00
Alex Skryl	f2f2d16451	Export: preserve Dtype state values in export callback arguments (#3145 ) Co-authored-by: Awni Hannun <awni@apple.com>	2026-02-19 08:07:28 -08:00
Awni Hannun	daf18e76ca	Fix fence synchronization accross command buffers (#3144 )	2026-02-18 20:21:26 -08:00
Anastasiia Filippova	06305022ab	Tensor scale nvfp4 (#3022 )	2026-02-18 11:19:26 +01:00
willem adnet	360639c2df	Add the hamming window function (#3135 )	2026-02-17 00:56:05 -08:00
willem adnet	3bbe87e6dc	Add hanning window function (#3124 )	2026-02-16 09:44:49 -08:00
vskiwi	e226af720e	Propagate quantization mode in quantized layers (#3133 )	2026-02-15 18:33:13 -08:00
Awni Hannun	43f4a74826	Manage stream placement in import function (#3127 )	2026-02-15 06:17:06 -08:00
Angelos Katharopoulos	c184262d29	Fix donation in sdpa vector (#3121 )	2026-02-12 10:46:21 -08:00
Cheng	72e94c81e1	[CUDA] Attention sinks in cuDNN SDPA (#3118 )	2026-02-11 16:46:39 +09:00
Awni Hannun	4c86c1e55a	Fix precision in Metal fused attention (#3119 )	2026-02-10 14:18:29 -08:00
Anastasiia Filippova	be52cf660b	register pressure (#3116 )	2026-02-10 14:17:28 -08:00
Cheng	54bb3eea42	[CUDA] Use cuDNN SDPA for decoding when using fixed-size KV cache (#3113 )	2026-02-10 09:15:45 +09:00
Anastasiia Filippova	5e018de4e5	Quantize module to QQLinear (#3106 )	2026-02-09 14:35:17 -08:00
Cheng	9cd4b9be91	[CUDA] Set current device before allocating memory (#3110 )	2026-02-08 19:04:57 +09:00
Cheng	566bc16b7c	Cleanup test_fast_sdpa.py (#3112 )	2026-02-08 19:04:24 +09:00
Awni Hannun	8fe1d09207	Fix residency set with user provided buffer (#3108 )	2026-02-06 16:38:36 -08:00
Ronan Collobert	ef3fbc60a3	is_available() should check the device index too (#3107 )	2026-02-06 13:02:04 -08:00
Angelos Katharopoulos	69fd3fa9b1	Patch bump (#3102 )	2026-02-06 09:15:22 -08:00
Awni Hannun	185b06d9ef	Patch for multi device CUDA (#3100 )	2026-02-05 17:33:51 -08:00
Manuel Candales	90e38f7b93	Fix qmv_impl for small N (#3096 )	2026-02-05 17:33:36 -08:00
Angelos Katharopoulos	ceea571490	JACCL update (#3094 )	2026-02-05 15:16:07 -08:00
Awni Hannun	99ca62c4d3	Fix 2pass sdpa on < M2 (#3099 )	2026-02-05 08:51:29 -08:00
Awni Hannun	206cf07e5b	Fix non simd f16 build (#3097 )	2026-02-05 07:04:02 -08:00
Jesse Gross	f47729c0d8	Disable managed memory on WSL when concurrentManagedAccess is not supported (#3095 )	2026-02-05 10:58:49 +09:00
Awni Hannun	b9b672250e	patch (#3093 )	2026-02-03 07:24:30 -08:00
Awni Hannun	adcbb91a9e	Fix for NAX overflow. (#3092 )	2026-02-02 18:54:01 -08:00
Awni Hannun	b56782be52	[Metal] Tune splitk gemm dispatch conditions and partition sizes (#3087 )	2026-02-02 08:45:09 -08:00
Cheng	8ef539522c	Fix failing python tests on Windows (#3076 )	2026-01-30 17:50:18 +09:00
Cheng	212077f163	Fallback to pinned host memory when managed memory is not supported (#3075 )	2026-01-30 13:18:41 +09:00
Awni Hannun	cc6e4eebad	Fix nax condition for iphone (#3083 )	2026-01-29 13:30:37 -08:00
Awni Hannun	fcbdd05022	More useful error for large indices (#3079 )	2026-01-29 13:02:39 -08:00
atharva	590b4f1c16	Fix ALiBi slopes for non-power-of-2 num_heads (#3071 )	2026-01-29 07:23:11 -08:00
Anri Lombard	0c6a895ed7	Use lower-right causal mask alignment consistently (#2967 ) Co-authored-by: Awni Hannun <awni.hannun@gmail.com>	2026-01-28 17:15:14 -08:00
stef	c86a9bced1	[Docs] Simple example of using MLX distributed (#2973 ) Co-authored-by: Awni Hannun <awni@apple.com>	2026-01-28 17:14:56 -08:00
Awni Hannun	12e386f308	Tune CUDA gaph sizes on B200 and H100 (#3077 )	2026-01-28 17:14:44 -08:00
Cheng	2ac18eddb9	[CUDA] Fallback Event impl when there is no hardware cpu/gpu coherency (#3070 )	2026-01-28 10:43:22 +09:00
Awni Hannun	b537b3685f	patch (#3074 )	2026-01-28 10:42:55 +09:00
Awni Hannun	2f324cc3b2	remove thrust (#3067 )	2026-01-27 08:54:07 -08:00
Awni Hannun	4912cc47c2	Fp qmv (#2984 )	2026-01-27 06:33:06 -08:00
Cheng	ce4d0a62ef	Do not require ConcurrentManagedAccess when not used (#3062 )	2026-01-27 11:19:20 +09:00
Cheng	73136472e0	Delay load CUDA libs and resolve DLL paths at runtime (#3061 )	2026-01-27 11:01:58 +09:00
Jesse Gross	fed0fe3c73	Better support consumer CUDA GPUs (#3056 )	2026-01-26 16:45:02 -08:00
Cheng	343ddf0d73	Fix long cache file path on Windows (#3065 )	2026-01-27 08:53:26 +09:00
Daniel Hiltgen	b70fc33ada	Improve CPU discovery (#3068 )	2026-01-26 15:01:43 -08:00
Merlin78	7ed2b6b935	Add NAX Split-K GEMM for large-K matmuls to improve performance (#3018 ) Co-authored-by: Huan <huan_xu@apple.com>	2026-01-26 11:23:20 -08:00
Daniel Hiltgen	a828e769be	GPU discovery (#3055 ) Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2026-01-26 09:54:13 -08:00
Nripesh Niketan	b6aa03e5b8	Update pre-commit hooks and versions for clang-format, black, and isort (#3059 )	2026-01-26 06:57:04 -08:00
Awni Hannun	5bd99dd5ec	Fix flaky macOS test (#3063 )	2026-01-25 16:40:57 -08:00
Awni Hannun	9e2d2a5957	[CUDA] Fast sorting (#3060 )	2026-01-25 15:10:18 -08:00
Cheng	3ac892b008	Hide symbols by default for mac/linux (#3057 )	2026-01-25 14:30:41 +09:00
Cheng	0bb50d99c0	Fix some NVCC warnings when building CUDA backend with MSVC (#3038 )	2026-01-25 12:25:01 +09:00
Cheng	257c422a8c	Find system-installed cuDNN on Windows (#3052 )	2026-01-25 12:24:22 +09:00
Awni Hannun	1935ab4452	Faster two pass sdpa (#3023 )	2026-01-24 14:16:33 -08:00
Cheng	617fd9cbbd	Use C++20 (#3050 )	2026-01-24 08:48:41 +09:00
Cheng	8e93b7448c	Fix some MSVC compilation errors (#3048 )	2026-01-24 07:56:56 +09:00
Cheng	fd27829efa	Build and test python package on Windows CI (#3049 )	2026-01-24 07:22:36 +09:00
Anri Lombard	dc81c1503a	Add missing <algorithm> include to buffer_cache.h (#3053 )	2026-01-23 11:52:36 -08:00
Awni Hannun	9bac6f8584	Allow take on empty array when it makes sense (#3046 )	2026-01-23 07:25:46 -08:00
Cheng	1650c4905a	Link with prebuilt OpenBLAS and fix shared libs build on Windows (#3036 )	2026-01-23 11:17:26 +09:00
Angelos Katharopoulos	becc769012	CUDA gather mv (#3039 )	2026-01-22 17:20:48 -08:00
Daniel Hiltgen	687508dd98	win: symbol exports and minor fixes (#3024 ) Co-authored-by: Cheng <zcbenz@gmail.com>	2026-01-23 10:16:22 +09:00
Cheng	c46c3833ee	Use cuda::std for math ops (#3041 )	2026-01-23 08:38:26 +09:00
Cheng	faea3e6d34	Turn nccl_stub into a normal target (#3037 )	2026-01-23 08:12:31 +09:00
Anastasiia Filippova	d98776e190	Columnwise quantize (#2989 )	2026-01-22 06:08:56 -08:00
Cheng	b2f86214bb	Remove xmlrunner from macOS CI (#3032 )	2026-01-22 08:06:28 +09:00
Awni Hannun	f28f9f0155	build 26.0 release in actions (#3035 )	2026-01-21 14:04:14 -08:00
rltakashige	0d698bc9a5	Handle data smaller than BUFFER_SIZE in jaccl recv (#3033 )	2026-01-21 13:44:41 -08:00
Awni Hannun	1d56dfdf59	Use higher precision for linspace with double (#3029 )	2026-01-21 06:20:50 -08:00
Dan Anderson	9a277a277a	PR 3007 Fix Seg Fault (#3008 ) Co-authored-by: KD2YCU <me@kd2ycu.com>	2026-01-20 21:39:15 -08:00
Cheng	8017d438a9	[CUDA] Faster grouped mm (#3011 )	2026-01-21 09:30:12 +09:00
Robert	634b148dd4	Optimize erf function with expm1f in Metal backend (#3025 )	2026-01-20 15:57:12 -08:00
Cheng	bfd62a50f4	Windows CI (#3021 )	2026-01-21 08:06:32 +09:00
Dan Anderson	83bb7891db	Fix negative dim indexing (#2994 ) Co-authored-by: KD2YCU <me@kd2ycu.com> Co-authored-by: Awni Hannun <awni@apple.com>	2026-01-20 06:24:33 -08:00
Cheng	65b42c8476	Do not give workflow boolean inputs default values (#3014 )	2026-01-20 15:27:14 +09:00
Cheng	0b25c9c06c	Do not clear disk space in setup-linux (#3013 )	2026-01-20 07:22:19 +09:00
XXXXRT666	46d0fdc5ec	Type Enhancement for Func Transforms and Bug Fix (#3003 ) Co-authored-by: Awni Hannun <awni@apple.com>	2026-01-20 07:19:57 +09:00
Cheng	d96a2bdf57	Fix python package install path in stubgen (#3009 )	2026-01-19 09:34:02 +09:00
Cheng	9052f678b3	Update CCCL to v3.1.3 (#3012 )	2026-01-19 07:50:09 +09:00
Tarjei Mandt	ca14d3d835	Fix sharding of quantized models with non-power-of-2 bits (#3006 )	2026-01-18 07:21:56 -08:00
gufengc	d2bef3c6bb	fix distributed all_to_sharded bias shard axis from -2 to -1 (#2987 )	2026-01-17 06:51:42 -08:00
Angelos Katharopoulos	3fe7794f22	Reverts changing the MLX_IBV_DEVICES to MLX_JACCL_DEVICES (#2999 )	2026-01-14 15:44:17 -08:00
Awni Hannun	47430159fc	Fix fence (#2998 )	2026-01-14 11:59:09 -08:00
Awni Hannun	2469fc2939	patch bump for next release (#2991 )	2026-01-14 08:46:09 -08:00
Awni Hannun	ac26a4cc0d	Allow some non 2D inputs in qqmm (#2981 )	2026-01-13 15:48:30 -08:00
Awni Hannun	099dcc0f4c	Expose to/from fp8 in Python and don't auto-convert fp8 when loading from safetensors (#2985 )	2026-01-13 15:48:21 -08:00
Awni Hannun	8654b8281d	Don't try to use NAX at run-time if kernels aren't there (#2982 )	2026-01-13 15:47:45 -08:00
MillaFleurs	4160ec10f7	Fix RandomBits::is_equivalent to include width (#2978 ) Co-authored-by: KD2YCU <me@kd2ycu.com> Co-authored-by: Angelos Katharopoulos <katharas@gmail.com> Co-authored-by: Awni Hannun <awni@apple.com>	2026-01-13 12:42:37 -08:00
Evan Quiney	a8197795f5	replace MLX_IBV_COORDINATOR with MLX_JACCL_COORDINATOR (#2986 )	2026-01-13 11:26:25 -08:00
CCYeh	7b1c46982a	fix doc (#2988 )	2026-01-12 13:33:26 -08:00
Anri Lombard	edab937248	Add asarray to __array_namespace__ (#2966 )	2026-01-12 06:16:27 -08:00
CCYeh	46ee0e9068	Fix `grid_dim_x` calculations (#2980 )	2026-01-12 06:16:05 -08:00
Anastasiia Filippova	43341e8d53	Swizzle scales (#2979 )	2026-01-10 15:32:54 -08:00