Fix for NAX overflow. (#3092 )

[Metal] Tune splitk gemm dispatch conditions and partition sizes (#3087 )
Fix failing python tests on Windows (#3076 )
2026-02-02 18:54:01 -08:00 · 2026-02-02 08:45:09 -08:00 · 2026-01-30 17:50:18 +09:00 · 2026-01-30 13:18:41 +09:00 · 2026-01-29 13:30:37 -08:00 · 2026-01-29 13:02:39 -08:00
236 changed files with 7019 additions and 2429 deletions
@@ -9,6 +9,7 @@ inputs:
 runs:
  using: "composite"
  steps:
+
    - name: Install Python package
      id: python_build
      shell: sh
@@ -24,6 +25,8 @@ runs:
          # 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'
@@ -35,6 +35,7 @@ jobs:
          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'
@@ -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
@@ -110,6 +110,11 @@ 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
        with:
@@ -161,7 +166,7 @@ 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
@@ -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,7 +194,7 @@ 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
@@ -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,7 +217,7 @@ 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
@@ -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,7 +240,7 @@ 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
@@ -245,7 +250,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/
@@ -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,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,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,7 @@ Build from source
 Build Requirements
 ^^^^^^^^^^^^^^^^^^

- A C++ compiler with C++17 support (e.g. Clang >= 5.0)
+- 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
@@ -183,3 +183,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:

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

@@ -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>
@@ -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;
 }
@@ -14,6 +14,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,
@@ -922,20 +935,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 +946,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 +1054,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 +1237,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 +1246,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 +1328,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;
@@ -105,7 +112,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 +180,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,12 @@ 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;
+    auto converted = *(Simd<float16_t, N>*)(&v);
+    converted = converted * 256.0;
+    auto sign = Simd<bool, N>(x & 128);
+    Simd<float, N> out = select(sign, -converted, converted);
+    return 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
@@ -56,7 +56,10 @@ target_sources(
          ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/affine_quantize.cu
          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/fp_quantize.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qmv.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)

@@ -66,12 +69,12 @@ add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/unary)
 # fp4 is not available on < 12.8
 if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 12.8.0)
  target_include_directories(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/)
+  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 +116,16 @@ 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)
-  target_compile_options(
-    mlx
-    PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--static-global-template-stub=false>")
-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.
@@ -157,16 +154,46 @@ 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")
+# 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 +201,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 +229,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 +251,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,11 +3,13 @@
 #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>

@@ -20,6 +22,49 @@ 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;
+
+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;
+      }
+#if defined(_WIN32)
+      // Empirically on Windows if there is no concurrentManagedAccess the
+      // managed memory also does not work.
+      if (!d.concurrent_managed_access()) {
+        return false;
+      }
+#endif
+    }
+    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 +78,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 +104,7 @@ SmallSizePool::SmallSizePool() {
 }

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

@@ -96,39 +138,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 +163,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 +175,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 +194,13 @@ 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));
+        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 +216,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 +235,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 +262,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 +274,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 +359,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 +388,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_;
 };
@@ -346,11 +346,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())));
      }
    });
  });
@@ -376,11 +376,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));
      }
    }

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

@@ -98,7 +98,7 @@ 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;

@@ -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;
 }
@@ -124,13 +125,14 @@ class DnnGraph : public fe::graph::Graph {

  // 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.
@@ -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) 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>
@@ -37,30 +38,25 @@ 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() {
@@ -81,6 +77,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 +214,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;
@@ -409,14 +419,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 +496,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>

@@ -119,7 +118,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 +141,7 @@ class Device {
  explicit Device(int device);
  ~Device();

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

@@ -149,6 +149,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 +161,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.
 ///////////////////////////////////////////////////////////////////////////////
@@ -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);
@@ -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,94 @@ 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();
+}
+
+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;
 }

 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;
+  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
+  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(uint64_t value) {
+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 +273,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 +301,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
@@ -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();

-  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
@@ -29,15 +29,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)));
  }
 }

@@ -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));
 }
@@ -201,4 +263,65 @@ 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,
+          0,
+          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
@@ -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 {
@@ -96,7 +97,78 @@ __global__ void prepare_grouped_mm_data(

 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(
+        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 +181,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 +280,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,9 +302,9 @@ 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);
@@ -229,8 +317,8 @@ void cutlass_grouped_gemm_unaligned(
      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 +337,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,
@@ -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));
@@ -179,14 +180,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));
@@ -258,14 +260,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"));
        }
      }
    }
@@ -360,15 +363,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"));
        }
      }
    }
@@ -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");
  }
@@ -272,7 +281,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 +330,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 +340,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 +393,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);
@@ -123,19 +123,19 @@ std::pair<dim3, dim3> get_grid_and_block(int dim0, int dim1, int dim2);

 // Get the num_blocks and block_dims assuming each thread handles
 // |work_per_thread| elements of |arr|.
-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);

-inline std::tuple<dim3, uint> get_launch_args(
+inline std::tuple<dim3, uint32_t> get_launch_args(
    const array& arr,
    bool large,
    int work_per_thread = 1,
-    uint max_block_dim = 1024) {
+    uint32_t max_block_dim = 1024) {
  return get_launch_args(
      arr.size(),
      arr.shape(),
@@ -89,11 +89,12 @@ class LRUCache {
    }

    if (env_name_ && ++cache_misses_ > 2 * capacity_) {
-      throw std::runtime_error(fmt::format(
-          "Cache thrashing is happening, please set the environment variable "
-          "{} to a larger value than {} to fix degraded performance.",
-          env_name_,
-          capacity_));
+      throw std::runtime_error(
+          fmt::format(
+              "Cache thrashing is happening, please set the environment variable "
+              "{} to a larger value than {} to fix degraded performance.",
+              env_name_,
+              capacity_));
    }

    vlist_.emplace_front(key, std::forward<U>(value));
@@ -354,6 +354,19 @@ void GatherMM::eval_gpu(const std::vector<array>& inputs, array& out) {
    return;
  }

+  auto [transposed_a, lda, a_] = check_transpose(encoder, s, a);
+  auto [transposed_b, ldb, b_] = check_transpose(encoder, s, b);
+  auto use_gemv = cu::can_use_gemv(M, N, K, transposed_a, transposed_b);
+  if (M == 1 && use_gemv) {
+    gather_mv(b_, a_, rhs_indices, lhs_indices, out, N, K, encoder);
+    return;
+  }
+
+  if (N == 1 && use_gemv) {
+    gather_mv(a_, b_, lhs_indices, rhs_indices, out, M, K, encoder);
+    return;
+  }
+
  throw std::runtime_error("NYI");
 }

@@ -24,20 +24,12 @@ namespace mlx::core {
    throw std::runtime_error(#func " has no CUDA implementation.");   \
  }

-#if CUDART_VERSION < 12080
-void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
-  throw std::runtime_error(
-      "[QQMatmul::eval_gpu] QQMM is only supported with CUDA 12.8 or higher.");
-}
-#endif
-
 NO_GPU(BlockMaskedMM)
 NO_GPU(FFT)
 NO_GPU(GatherQMM)
 NO_GPU(Hadamard)
 NO_GPU_MULTI(LUF)
 NO_GPU_MULTI(QRF)
-NO_GPU(QuantizedMatmul)
 NO_GPU(SegmentedMM)
 NO_GPU_MULTI(SVD)
 NO_GPU(Inverse)
@@ -23,8 +23,7 @@ affine_quantize(const T* w, uint8_t* out, T* scales, T* biases, size_t size) {
  auto tidx = block_idx.x * block_size.x + idx_in_block.x;
  auto tidy = block_idx.y * block_size.y + idx_in_block.y;

-  auto grid_dim_x =
-      cg::this_grid().dim_blocks().x * cg::this_grid().block_index().x;
+  auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x;
  constexpr float eps = 1e-7;
  constexpr int simd_size = WARP_SIZE;
  constexpr float n_bins = (1 << bits) - 1;
@@ -141,8 +140,7 @@ __global__ void affine_dequantize(
  auto tidx = block_idx.x * block_size.x + idx_in_block.x;
  auto tidy = block_idx.y * block_size.y + idx_in_block.y;

-  auto grid_dim_x =
-      cg::this_grid().dim_blocks().x * cg::this_grid().block_index().x;
+  auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x;

  constexpr int pack_factor = get_pack_factor<bits, 8>();
  constexpr int bytes_per_pack = get_bytes_per_pack<bits>();
@@ -210,7 +208,7 @@ __global__ void affine_dequantize(
        bias;
    out[3] = static_cast<T>((w[2] >> 2) & 0x3f) * scale + bias;
  } else {
-    uint val = w[offset];
+    uint32_t val = w[offset];
 #pragma clang loop unroll(full)
    for (int i = 0; i < pack_factor; i++) {
      uint8_t d;
@@ -81,3 +81,20 @@ struct __nv_fp4_e2m1 {
  }
  uint8_t __x{0};
 };
+
+struct __nv_fp4x4_e2m1 {
+  __device__ operator float4() {
+    float4 out;
+    auto bits = __high & 0xf;
+    out.x = float(*(__nv_fp4_e2m1*)(&bits));
+    bits = (__high >> 4) & 0xf;
+    out.y = float(*(__nv_fp4_e2m1*)(&bits));
+    bits = (__low) & 0xf;
+    out.z = float(*(__nv_fp4_e2m1*)(&bits));
+    bits = (__low >> 4) & 0xf;
+    out.w = float(*(__nv_fp4_e2m1*)(&bits));
+    return out;
+  }
+  uint8_t __high{0};
+  uint8_t __low{0};
+};
@@ -11,8 +11,6 @@

 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
-#include <cuda_fp4.h>
-#include <cuda_fp8.h>

 namespace mlx::core {
 namespace cu {
@@ -31,7 +29,71 @@ struct Dequantize {
 namespace cg = cooperative_groups;

 template <typename T, int group_size, int bits, bool use_mx_scale, bool USE_SR>
-__global__ void fp_quantize(T* w, uint8_t* out, uint8_t* scales, size_t size) {
+__global__ void fp_quantize_dequantize(T* w, T* out, size_t size) {
+  using Tx2 = Vector2_t<T>;
+  using Tx4 = Vector4_t<T>;
+  uint32_t rbits = 0; // reserved bits for future use
+  auto block_size = cg::this_thread_block().dim_threads();
+  auto block_idx = cg::this_thread_block().group_index();
+  auto idx_in_block = cg::this_thread_block().thread_index();
+  auto tidx = block_idx.x * block_size.x + idx_in_block.x;
+  auto tidy = block_idx.y * block_size.y + idx_in_block.y;
+  auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x;
+
+  size_t thread_idx = tidx + grid_dim_x * size_t(tidy);
+  size_t base_idx = thread_idx * group_size;
+
+  if (base_idx >= size) {
+    return;
+  }
+
+  auto w_tile = load_vector<group_size, T>(w, thread_idx);
+  float scale = 0.0f;
+
+  Tx2 amax_2x = Tx2{0.0f, 0.0f};
+
+#pragma unroll
+  for (int i = 0; i < group_size; i += 2) {
+    auto pair = Tx2{w_tile[i], w_tile[i + 1]};
+    abs_max_x2<Tx2>(amax_2x, amax_2x, pair);
+  }
+
+  scale = static_cast<float>(
+      max(fabsf(static_cast<float>(amax_2x.x)),
+          fabsf(static_cast<float>(amax_2x.y))));
+
+  scale /= bits == 4 ? 6.0f : 448.0f;
+  // Convert to mx scale or nv scale
+  using ScaleType =
+      std::conditional_t<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
+  auto s = ScaleType(scale);
+  scale = float(s);
+  AlignedVector<T, group_size> w_hat;
+
+#pragma unroll
+  for (int i = 0; i < group_size / 4; i++) {
+    Tx4 w_Tx4 = *reinterpret_cast<Tx4*>(&w_tile[i * 4]);
+    float4 dq;
+    if constexpr (bits == 8) {
+      uint32_t quantized_val =
+          scale_cvt_Tx4_to_fp8x4<T, USE_SR>(w_Tx4, 1.0f / scale, rbits);
+      dq = dequant_fp8(quantized_val);
+    } else {
+      uint16_t quantized_val =
+          scale_cvt_Tx4_to_fp4x4<T, USE_SR>(w_Tx4, 1.0f / scale, rbits);
+      dq = dequant_fp4(quantized_val);
+    }
+    w_hat[i * 4] = static_cast<T>(dq.x * scale);
+    w_hat[i * 4 + 1] = static_cast<T>(dq.y * scale);
+    w_hat[i * 4 + 2] = static_cast<T>(dq.z * scale);
+    w_hat[i * 4 + 3] = static_cast<T>(dq.w * scale);
+  }
+  store_vector<group_size>(out, thread_idx, w_hat);
+}
+
+template <typename T, int group_size, int bits, bool use_mx_scale, bool USE_SR>
+__global__ void
+fp_quantize_rowwise(T* w, uint8_t* out, uint8_t* scales, size_t size) {
  using Tx2 = Vector2_t<T>;
  using Tx4 = Vector4_t<T>;
  uint32_t rbits = 0; // reserved bits for future use
@@ -92,6 +154,133 @@ __global__ void fp_quantize(T* w, uint8_t* out, uint8_t* scales, size_t size) {
  store_vector<group_size / elem_per_byte>(out, thread_idx, quantized);
 }

+template <typename T, int group_size, int bits, bool use_mx_scale, bool USE_SR>
+__global__ void fp_quantize_columnwise(
+    T* w,
+    uint8_t* out,
+    uint8_t* scales,
+    size_t size,
+    int M,
+    int K) {
+  // Input: [M, K] with strides [1, M] (M-major)
+  // Quantized output: [M, K/elem_per_byte] row-major (K-major)
+  // Scales: [M, K/group_size] row-major (K-major)
+  // Quantize along K (last dimension, groups of group_size elements)
+
+  using Tx2 = Vector2_t<T>;
+  using Tx4 = Vector4_t<T>;
+  uint32_t rbits = 0;
+
+  auto block_idx = cg::this_thread_block().group_index();
+  auto idx_in_block = cg::this_thread_block().thread_index();
+
+  constexpr int BLOCK_X = 32;
+  constexpr int BLOCK_Y = 32;
+  constexpr int elem_per_byte = (bits == 8) ? 1 : 2;
+  constexpr int bytes_per_group = group_size / elem_per_byte;
+
+  constexpr int rows_per_block = BLOCK_X;
+  constexpr int cols_per_block = BLOCK_Y * group_size;
+  constexpr int local_cols = cols_per_block / elem_per_byte;
+  constexpr int bytes_per_block = rows_per_block * local_cols;
+
+  constexpr int SMEM_PAD = 4;
+  constexpr int padded_local_cols = local_cols + SMEM_PAD;
+
+  auto tidx = idx_in_block.x;
+  auto tidy = idx_in_block.y;
+
+  int num_col_blocks = (K + cols_per_block - 1) / cols_per_block;
+  auto bidx = block_idx.x % num_col_blocks;
+  auto bidy = block_idx.x / num_col_blocks;
+
+  T thread_data[group_size];
+
+  __shared__ uint8_t quantized_smem[rows_per_block * padded_local_cols];
+  __shared__ uint8_t scales_smem[BLOCK_X][BLOCK_Y + SMEM_PAD];
+
+  int row_base = bidy * rows_per_block + tidx;
+  int col_base = bidx * cols_per_block + tidy * group_size;
+
+  bool valid = (row_base < M) && (col_base + group_size <= K);
+  if (valid) {
+#pragma unroll
+    for (int i = 0; i < group_size; i++) {
+      auto index = row_base + (col_base + i) * M;
+      thread_data[i] = w[index];
+    }
+
+    // Compute scale
+    Tx2 amax_2x = Tx2{0.0f, 0.0f};
+#pragma unroll
+    for (int r = 0; r < group_size; r += 2) {
+      auto pair = Tx2{thread_data[r], thread_data[r + 1]};
+      abs_max_x2<Tx2>(amax_2x, amax_2x, pair);
+    }
+    float scale =
+        max(fabsf(static_cast<float>(amax_2x.x)),
+            fabsf(static_cast<float>(amax_2x.y)));
+    scale /= (bits == 4) ? 6.0f : 448.0f;
+    using ScaleType =
+        std::conditional_t<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
+    auto s = ScaleType(scale);
+    scale = float(s);
+    scales_smem[tidx][tidy] = s.__x;
+
+    int shared_idx = tidx * padded_local_cols + tidy * bytes_per_group;
+
+#pragma unroll
+    for (int j = 0; j < group_size / 4; j++) {
+      Tx4 w_Tx4 = *reinterpret_cast<Tx4*>(&thread_data[j * 4]);
+      if constexpr (bits == 8) {
+        uint32_t quantized_val =
+            scale_cvt_Tx4_to_fp8x4<T, USE_SR>(w_Tx4, 1.0f / scale, rbits);
+        *reinterpret_cast<uint32_t*>(&quantized_smem[shared_idx + j * 4]) =
+            quantized_val;
+      } else {
+        uint16_t quantized_val =
+            scale_cvt_Tx4_to_fp4x4<T, USE_SR>(w_Tx4, 1.0f / scale, rbits);
+        *reinterpret_cast<uint16_t*>(&quantized_smem[shared_idx + j * 2]) =
+            quantized_val;
+      }
+    }
+  }
+  __syncthreads();
+
+  int output_cols = K / elem_per_byte;
+  int num_groups_per_row = K / group_size;
+  int linear_tid = tidx + tidy * BLOCK_X;
+  // Write back quantized values
+#pragma unroll
+  for (int i = linear_tid; i < bytes_per_block; i += BLOCK_X * BLOCK_Y) {
+    int local_row = i / local_cols;
+    int local_col = i % local_cols;
+
+    int global_row = bidy * rows_per_block + local_row;
+    int global_col = bidx * local_cols + local_col;
+
+    if (global_row < M && global_col < output_cols) {
+      int physical_idx = local_row * padded_local_cols + local_col;
+      out[global_row * output_cols + global_col] = quantized_smem[physical_idx];
+    }
+  }
+  // Write back scales
+  constexpr int num_scales = BLOCK_X * BLOCK_Y;
+#pragma unroll
+  for (int i = linear_tid; i < num_scales; i += BLOCK_X * BLOCK_Y) {
+    int local_row = i / BLOCK_Y;
+    int local_col = i % BLOCK_Y;
+
+    int global_row = bidy * BLOCK_X + local_row;
+    int global_col = bidx * BLOCK_Y + local_col;
+
+    if (global_row < M && global_col < num_groups_per_row) {
+      scales[global_row * num_groups_per_row + global_col] =
+          scales_smem[local_row][local_col];
+    }
+  }
+}
+
 template <typename T, int group_size, int bits, bool use_mx_scale>
 __global__ void
 fp_dequantize(const uint8_t* w, const uint8_t* scales, T* out, size_t size) {
@@ -119,7 +308,7 @@ fp_dequantize(const uint8_t* w, const uint8_t* scales, T* out, size_t size) {

  out += oindex;

-  uint val = w[offset];
+  uint32_t val = w[offset];
 #pragma clang loop unroll(full)
  for (int i = 0; i < pack_factor; i++) {
    uint8_t d;
@@ -132,8 +321,60 @@ fp_dequantize(const uint8_t* w, const uint8_t* scales, T* out, size_t size) {
  }
 }

+inline std::tuple<dim3, dim3>
+get_columnwise_quantize_launch_args(size_t size, int group_size, int M, int K) {
+  constexpr int BLOCK_X = 32;
+  constexpr int BLOCK_Y = 32;
+  int rows_per_block = BLOCK_X;
+  int cols_per_block = BLOCK_Y * group_size;
+
+  dim3 grid;
+  grid.x =
+      cuda::ceil_div(K, cols_per_block) * cuda::ceil_div(M, rows_per_block);
+  grid.y = 1;
+  grid.z = 1;
+
+  dim3 block(BLOCK_X, BLOCK_Y);
+
+  return std::make_tuple(grid, block);
+}
+
 } // namespace cu

+void fp_quantize_dequantize(
+    const array& w,
+    array& what,
+    int group_size,
+    int bits,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  enc.set_input_array(w);
+  enc.set_output_array(what);
+  dispatch_float_types(w.dtype(), "fp_quantize_dequantize", [&](auto type_tag) {
+    using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    if constexpr (!std::is_same_v<T, double>) {
+      auto kernel = cu::fp_quantize_dequantize<T, 32, 4, true, false>;
+      if (bits == 8) {
+        kernel = cu::fp_quantize_dequantize<T, 32, 8, true, false>;
+      } else if (group_size == 16) {
+        kernel = cu::fp_quantize_dequantize<T, 16, 4, false, false>;
+      }
+      bool large = w.size() > UINT_MAX;
+      auto [num_blocks, block_dims] =
+          get_launch_args(w.size(), w.shape(), w.strides(), large, group_size);
+
+      enc.add_kernel_node(
+          kernel,
+          num_blocks,
+          block_dims,
+          0,
+          gpu_ptr<T>(w),
+          gpu_ptr<T>(what),
+          w.size());
+    }
+  });
+}
+
 void fp_quantize(
    const array& w,
    array& wq,
@@ -145,33 +386,65 @@ void fp_quantize(
  enc.set_input_array(w);
  enc.set_output_array(wq);
  enc.set_output_array(scales);
-  dispatch_float_types(w.dtype(), "fp_quantize", [&](auto type_tag) {
-    using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
-    if constexpr (!std::is_same_v<T, double>) {
-      auto kernel = cu::fp_quantize<T, 32, 4, true, false>;
-      if (bits == 8) {
-        kernel = cu::fp_quantize<T, 32, 8, true, false>;
-      } else if (group_size == 16) {
-        kernel = cu::fp_quantize<T, 16, 4, false, false>;
+  if (w.strides().back() != 1) {
+    dispatch_float_types(w.dtype(), "fp_quantize_columnwise", [&](auto type_tag) {
+      using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+      if constexpr (!std::is_same_v<T, double>) {
+        auto M = w.shape(-2);
+        auto K = w.shape(-1);
+        auto kernel = cu::fp_quantize_columnwise<T, 32, 4, true, false>;
+        if (bits == 8) {
+          kernel = cu::fp_quantize_columnwise<T, 32, 8, true, false>;
+        } else if (group_size == 16) {
+          kernel = cu::fp_quantize_columnwise<T, 16, 4, false, false>;
+        }
+        auto [num_blocks, block_dims] =
+            cu::get_columnwise_quantize_launch_args(w.size(), group_size, M, K);
+        enc.add_kernel_node(
+            kernel,
+            num_blocks,
+            block_dims,
+            0,
+            gpu_ptr<T>(w),
+            gpu_ptr<uint8_t>(wq),
+            gpu_ptr<uint8_t>(scales),
+            w.size(),
+            M,
+            K);
+      } else {
+        throw std::runtime_error(
+            "[Quantize::eval_gpu] Can not quantize input with type float64.");
      }
-      bool large = w.size() > UINT_MAX;
-      auto [num_blocks, block_dims] =
-          get_launch_args(w.size(), w.shape(), w.strides(), large, group_size);
+    });
+  } else {
+    dispatch_float_types(w.dtype(), "fp_quantize_rowwise", [&](auto type_tag) {
+      using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+      if constexpr (!std::is_same_v<T, double>) {
+        auto kernel = cu::fp_quantize_rowwise<T, 32, 4, true, false>;
+        if (bits == 8) {
+          kernel = cu::fp_quantize_rowwise<T, 32, 8, true, false>;
+        } else if (group_size == 16) {
+          kernel = cu::fp_quantize_rowwise<T, 16, 4, false, false>;
+        }
+        bool large = w.size() > UINT_MAX;
+        auto [num_blocks, block_dims] = get_launch_args(
+            w.size(), w.shape(), w.strides(), large, group_size);

-      enc.add_kernel_node(
-          kernel,
-          num_blocks,
-          block_dims,
-          0,
-          gpu_ptr<T>(w),
-          gpu_ptr<uint8_t>(wq),
-          gpu_ptr<uint8_t>(scales),
-          w.size());
-    } else {
-      throw std::runtime_error(
-          "[Quantize::eval_gpu] Can not quantize input with type float64.");
-    }
-  });
+        enc.add_kernel_node(
+            kernel,
+            num_blocks,
+            block_dims,
+            0,
+            gpu_ptr<T>(w),
+            gpu_ptr<uint8_t>(wq),
+            gpu_ptr<uint8_t>(scales),
+            w.size());
+      } else {
+        throw std::runtime_error(
+            "[Quantize::eval_gpu] Can not quantize input with type float64.");
+      }
+    });
+  }
 }

 void fp_dequantize(
@@ -0,0 +1,26 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/quantized/qqmm_impl.h"
+
+namespace mlx::core {
+void qqmm_impl(
+    cu::CommandEncoder&,
+    int,
+    int,
+    int,
+    bool,
+    int64_t,
+    bool,
+    int64_t,
+    array&,
+    const array&,
+    const array&,
+    const array&,
+    const array&,
+    Dtype,
+    QuantizationMode,
+    float) {
+  throw std::runtime_error(
+      "[QQMatmul::eval_gpu] QQMM is only supported with CUDA 12.8 or higher.");
+}
+} // namespace mlx::core
@@ -0,0 +1,304 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device/utils.cuh"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/cuda/quantized/qmv.h"
+#include "mlx/backend/cuda/quantized/quantized_utils.cuh"
+#include "mlx/dtype_utils.h"
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+
+namespace mlx::core::cu {
+
+namespace cg = cooperative_groups;
+
+static constexpr int rows_per_block = 8;
+
+template <typename T>
+__device__ void adjust_matrix_offsets(
+    const T*& x,
+    const uint32_t*& w,
+    const uint8_t*& scales,
+    T*& y,
+    int output_stride,
+    const int& x_batch_ndims,
+    const Shape x_shape,
+    const Strides x_strides,
+    const int& w_batch_ndims,
+    const Shape w_shape,
+    const Strides w_strides,
+    const Strides s_strides) {
+  uint32_t idx = cg::this_grid().block_index().z;
+  if (x_batch_ndims == 1) {
+    x += idx * x_strides[0];
+  } else {
+    x += elem_to_loc(idx, x_shape.data(), x_strides.data(), x_batch_ndims);
+  }
+  if (w_batch_ndims == 1) {
+    w += idx * w_strides[0];
+    scales += idx * s_strides[0];
+  } else {
+    auto [w_idx, s_idx] = elem_to_loc(
+        idx, w_shape.data(), w_strides.data(), s_strides.data(), w_batch_ndims);
+    w += w_idx;
+    scales += s_idx;
+  }
+  y += idx * output_stride;
+}
+
+template <
+    typename T,
+    int rows_per_block,
+    int n_per_thread,
+    int bits,
+    int group_size,
+    bool use_mx_scale>
+__device__ void fp_qmv_impl(
+    const uint32_t* mat,
+    const uint8_t* scales_,
+    const T* vec,
+    T* out,
+    int rows,
+    int cols) {
+  auto block = cg::this_thread_block();
+  auto warp = cg::tiled_partition<WARP_SIZE>(block);
+
+  constexpr int vals_per_item = bits == 8 ? 4 : 8;
+  constexpr int nv_per_thread = vals_per_item * n_per_thread;
+  auto g_idx = block.group_index();
+  auto t_idx = block.thread_index();
+  int row = g_idx.y * rows_per_block + t_idx.y;
+
+  vec += g_idx.x * cols;
+  out += g_idx.x * rows;
+
+  using ScaleType =
+      std::conditional_t<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
+  auto scales = (ScaleType*)(scales_);
+  auto packed_cols = cols / vals_per_item;
+
+  if (row < rows) {
+    constexpr int scales_per_step = std::max(nv_per_thread / group_size, 1);
+    constexpr int scale_step = (WARP_SIZE * nv_per_thread) / group_size;
+    constexpr int n_per_step = n_per_thread / scales_per_step;
+    // Offset scales to correct row
+    scales += row * (cols / group_size) +
+        (warp.thread_rank() * nv_per_thread) / group_size;
+    float sum = 0.0f;
+    for (int col = n_per_thread * warp.thread_rank(); col < packed_cols;
+         col += (WARP_SIZE * n_per_thread)) {
+      auto local_vec =
+          unsafe_load_vector<nv_per_thread>(vec + vals_per_item * col, 0);
+      auto local_mat =
+          unsafe_load_vector<n_per_thread>(mat + row * packed_cols + col, 0);
+#pragma unroll
+      for (int i = 0; i < scales_per_step; ++i) {
+        float2 local_sum = {0.0f, 0.0f};
+#pragma unroll
+        for (int j = 0; j < n_per_step; ++j) {
+          int k = n_per_step * i + j;
+          if constexpr (bits == 8) {
+            auto v = dequant_fp8(local_mat[k]);
+            local_sum.x +=
+                v.x * static_cast<float>(local_vec[vals_per_item * k]);
+            local_sum.x +=
+                v.y * static_cast<float>(local_vec[vals_per_item * k + 1]);
+            local_sum.y +=
+                v.z * static_cast<float>(local_vec[vals_per_item * k + 2]);
+            local_sum.y +=
+                v.w * static_cast<float>(local_vec[vals_per_item * k + 3]);
+          } else {
+            auto v = dequant_fp4(local_mat[k]);
+            local_sum.x +=
+                v.x * static_cast<float>(local_vec[vals_per_item * k]);
+            local_sum.y +=
+                v.y * static_cast<float>(local_vec[vals_per_item * k + 1]);
+            local_sum.x +=
+                v.z * static_cast<float>(local_vec[vals_per_item * k + 2]);
+            local_sum.y +=
+                v.w * static_cast<float>(local_vec[vals_per_item * k + 3]);
+
+            v = dequant_fp4(local_mat[k] >> 16);
+            local_sum.x +=
+                v.x * static_cast<float>(local_vec[vals_per_item * k + 4]);
+            local_sum.y +=
+                v.y * static_cast<float>(local_vec[vals_per_item * k + 5]);
+            local_sum.x +=
+                v.z * static_cast<float>(local_vec[vals_per_item * k + 6]);
+            local_sum.y +=
+                v.w * static_cast<float>(local_vec[vals_per_item * k + 7]);
+          }
+        }
+        sum += (local_sum.x + local_sum.y) * float(scales[i]);
+      }
+      scales += scale_step;
+    }
+
+    sum = cg::reduce(warp, sum, cg::plus<float>{});
+    if (warp.thread_rank() == 0) {
+      out[row] = static_cast<T>(sum);
+    }
+  }
+}
+
+template <
+    typename T,
+    int rows_per_block,
+    int n_per_thread,
+    int bits,
+    int group_size,
+    bool use_mx_scale>
+__global__ void fp_qmv_single(
+    const uint32_t* mat,
+    const uint8_t* scales,
+    const T* vec,
+    T* out,
+    int rows,
+    int cols) {
+  fp_qmv_impl<T, rows_per_block, n_per_thread, bits, group_size, use_mx_scale>(
+      mat, scales, vec, out, rows, cols);
+}
+
+template <
+    typename T,
+    int rows_per_block,
+    int n_per_thread,
+    int bits,
+    int group_size,
+    bool use_mx_scale>
+__global__ void fp_qmv_batched(
+    const uint32_t* mat,
+    const uint8_t* scales,
+    const T* vec,
+    T* out,
+    int rows,
+    int cols,
+    int vec_batch_ndims,
+    const __grid_constant__ Shape vec_shape,
+    const __grid_constant__ Strides vec_strides,
+    int mat_batch_ndims,
+    const __grid_constant__ Shape mat_shape,
+    const __grid_constant__ Strides mat_strides,
+    const __grid_constant__ Strides scales_strides) {
+  adjust_matrix_offsets<T>(
+      vec,
+      mat,
+      scales,
+      out,
+      rows * vec_shape[vec_batch_ndims],
+      vec_batch_ndims,
+      vec_shape,
+      vec_strides,
+      mat_batch_ndims,
+      mat_shape,
+      mat_strides,
+      scales_strides);
+  fp_qmv_impl<T, rows_per_block, n_per_thread, bits, group_size, use_mx_scale>(
+      mat, scales, vec, out, rows, cols);
+}
+
+template <typename F>
+void dispatch_1_2_4(int n, F&& f) {
+  switch (n) {
+    case 1:
+      f(std::integral_constant<int, 1>{});
+      break;
+    case 2:
+      f(std::integral_constant<int, 2>{});
+      break;
+    case 4:
+      f(std::integral_constant<int, 4>{});
+      break;
+  }
+}
+
+void fp_qmv(
+    const array& mat,
+    const array& scales,
+    const array& vec,
+    array& out,
+    int bits,
+    int group_size,
+    int M,
+    int N,
+    int K,
+    CommandEncoder& encoder) {
+  encoder.set_input_array(mat);
+  encoder.set_input_array(scales);
+  encoder.set_input_array(vec);
+  encoder.set_output_array(out);
+  dispatch_float_types(out.dtype(), "qmv", [&](auto type_tag) {
+    using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    if constexpr (!std::is_same_v<T, double>) {
+      dim3 block_dims{WARP_SIZE, rows_per_block};
+      uint32_t B = out.size() / (M * N);
+      uint32_t blocks_y = (N + rows_per_block - 1) / rows_per_block;
+      const uint32_t* mat_ptr = gpu_ptr<uint32_t>(mat);
+      const T* vec_ptr = gpu_ptr<T>(vec);
+      int n = 1;
+      if (K % 32 == 0 && cu::is_aligned<4>(mat_ptr) &&
+          ((bits == 4 && cu::is_aligned<8>(vec_ptr)) ||
+           cu::is_aligned<4>(vec_ptr))) {
+        n = 4;
+      } else if (
+          cu::is_aligned<2>(mat_ptr) &&
+          ((bits == 4 && cu::is_aligned<4>(vec_ptr)) ||
+           cu::is_aligned<2>(vec_ptr))) {
+        n = 2;
+      }
+      dispatch_1_2_4(n, [&](auto n) {
+        dispatch_bool(B > 1, [&](auto batched) {
+          if (!batched.value) {
+            auto kernel =
+                fp_qmv_single<T, rows_per_block, n.value, 4, 32, true>;
+            if (bits == 8) {
+              kernel = fp_qmv_single<T, rows_per_block, n.value, 8, 32, true>;
+            } else if (group_size == 16) {
+              kernel = fp_qmv_single<T, rows_per_block, n.value, 4, 16, false>;
+            }
+            encoder.add_kernel_node(
+                kernel,
+                {static_cast<uint32_t>(M), blocks_y},
+                block_dims,
+                0,
+                mat_ptr,
+                gpu_ptr<uint8_t>(scales),
+                vec_ptr,
+                gpu_ptr<T>(out),
+                N,
+                K);
+          } else {
+            auto kernel =
+                fp_qmv_batched<T, rows_per_block, n.value, 4, 32, true>;
+            if (bits == 8) {
+              kernel = fp_qmv_batched<T, rows_per_block, n.value, 8, 32, true>;
+            } else if (group_size == 16) {
+              kernel = fp_qmv_batched<T, rows_per_block, n.value, 4, 16, false>;
+            }
+            encoder.add_kernel_node(
+                kernel,
+                {static_cast<uint32_t>(M), blocks_y, B},
+                block_dims,
+                0,
+                mat_ptr,
+                gpu_ptr<uint8_t>(scales),
+                vec_ptr,
+                gpu_ptr<T>(out),
+                N,
+                K,
+                vec.ndim() - 2,
+                const_param(vec.shape()),
+                const_param(vec.strides()),
+                mat.ndim() - 2,
+                const_param(mat.shape()),
+                const_param(mat.strides()),
+                const_param(scales.strides()));
+          }
+        });
+      });
+    }
+  });
+}
+
+} // namespace mlx::core::cu
@@ -0,0 +1,21 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device.h"
+
+namespace mlx::core::cu {
+
+void fp_qmv(
+    const array& w,
+    const array& scales,
+    const array& vec,
+    array& out,
+    int bits,
+    int group_size,
+    int M,
+    int N,
+    int K,
+    CommandEncoder& encoder);
+
+} // namespace mlx::core::cu
@@ -1,10 +1,11 @@
 // Copyright © 2025 Apple Inc.

 #include "mlx/backend/cuda/device.h"
-#include "mlx/backend/cuda/quantized/cublas_qqmm.h"
+#include "mlx/backend/cuda/quantized/qmv.h"
+#include "mlx/backend/cuda/quantized/qqmm_impl.h"
 #include "mlx/backend/cuda/quantized/qqmm_utils.h"
 #include "mlx/backend/cuda/quantized/quantized.h"
-#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/cuda/quantized/quantized_utils.h"
 #include "mlx/primitives.h"

 #include <nvtx3/nvtx3.hpp>
@@ -13,40 +14,7 @@ namespace mlx::core {

 namespace {

-inline array ensure_row_contiguous(
-    const array& x,
-    cu::CommandEncoder& enc,
-    const Stream& s) {
-  if (!x.flags().row_contiguous) {
-    array x_copy = contiguous_copy_gpu(x, s);
-    enc.add_temporary(x_copy);
-    return x_copy;
-  } else {
-    return x;
-  }
-}
-
-inline array ensure_row_contiguous_matrix(
-    const array& x,
-    cu::CommandEncoder& enc,
-    const Stream& s) {
-  if (x.ndim() < 2) {
-    if (x.strides()[0] == 1) {
-      return x;
-    }
-  } else {
-    auto stride_0 = x.strides()[x.ndim() - 2];
-    auto stride_1 = x.strides()[x.ndim() - 1];
-    if (stride_0 == x.shape(-1) && stride_1 == 1) {
-      return x;
-    }
-  }
-  array x_copy = contiguous_copy_gpu(x, s);
-  enc.add_temporary(x_copy);
-  return x_copy;
-}
-
-array pad_and_repack_scales(
+array pad_and_swizzle_scales(
    const array& scale,
    cu::CommandEncoder& encoder,
    const Stream& s) {
@@ -64,77 +32,63 @@ array pad_and_repack_scales(
      cu::malloc_async(pad_outer * pad_inner, encoder),
      Shape{pad_outer, pad_inner},
      scale.dtype());
-  repack_scales(scale, scale_tiled, encoder, s);
+  swizzle_scales(scale, scale_tiled, encoder, s);

  encoder.add_temporary(scale_tiled);
  return scale_tiled;
 }
-} // namespace

-namespace {
-void qqmm_impl(
-    cu::CommandEncoder& encoder,
-    int M,
-    int N,
-    int K,
-    bool a_transposed,
-    int64_t lda,
-    bool b_transposed,
-    int64_t ldb,
-    array& out,
-    const array& a,
-    const array& b,
-    const array& a_scale,
-    const array& b_scale,
-    Dtype out_dtype,
-    QuantizationMode mode,
-    float alpha = 1.0f) {
-  // Invoke CublasQQMM
-  std::string qmode = quantization_mode_to_string(mode);
-
-  // Currently only supports non-batched QQMM operations
-  // that covers all use cases for training, we will just collapse (batch,
-  // seq_len) into (tokens)
-  CublasQQMM qqmm(
-      encoder.device(),
-      a_transposed,
-      M,
-      K,
-      lda,
-      b_transposed,
-      K,
-      N,
-      ldb,
-      1, // batch_count
-      0, // a_batch_stride
-      0, // b_batch_stride
-      out_dtype,
-      qmode);
-
-  qqmm.run(encoder, out, a, b, a_scale, b_scale, alpha);
-}
 } // namespace

 void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(
+      (inputs.size() == 3 && inputs[1].dtype() == uint32) ||
+      (inputs.size() == 2));
  nvtx3::scoped_range r("QQMatmul::eval_gpu");
+
  auto& s = stream();
  auto& encoder = cu::get_command_encoder(s);
  auto& device = encoder.device();
+
+  bool w_quantized = (inputs[1].dtype() == uint32);
+  if (w_quantized && inputs[0].shape(-2) == 1) {
+    out.set_data(cu::malloc_async(out.nbytes(), encoder));
+
+    bool donate_x = inputs[0].is_donatable();
+    array x = ensure_row_contiguous(inputs[0], encoder, s);
+    // If x is a copy it should be donatable
+    donate_x |= x.is_donatable();
+    auto xhat = donate_x
+        ? x
+        : array(cu::malloc_async(x.nbytes(), encoder), x.shape(), x.dtype());
+    if (!donate_x) {
+      encoder.add_temporary(xhat);
+    }
+    fp_quantize_dequantize(x, xhat, group_size_, bits_, encoder, s);
+
+    // Make sure the last two dims of w and s are contiguous
+    array w = ensure_row_contiguous_matrix(inputs[1], encoder, s);
+    array scales = ensure_row_contiguous_matrix(inputs[2], encoder, s);
+
+    bool non_batched = w.ndim() == 2;
+    int K = x.shape(-1);
+    int M = non_batched ? x.size() / K : x.shape(-2);
+    int N = out.shape(-1);
+
+    fp_qmv(w, scales, xhat, out, bits_, group_size_, M, N, K, encoder);
+    return;
+  }
+
  auto cc = device.compute_capability_major() * 100 +
      device.compute_capability_minor() * 10;
  if (cc < 1000) {
    throw std::runtime_error(
        "[QQMatmul::eval_gpu] QQMM is only supported on GPUs with compute capability 10.0 or higher.");
  }
-  assert(
-      (inputs.size() == 3 && inputs[1].dtype() == uint32) ||
-      (inputs.size() == 2));
-
  auto quantize = [&](const array& input,
                      cu::CommandEncoder& encoder,
                      const Stream& s) -> std::pair<array, array> {
-    const array x = ensure_row_contiguous(input, encoder, s);
-
+    auto x = ensure_contiguous(input, encoder, s);
    auto xq_shape = x.shape();
    xq_shape.back() = x.shape(-1) * bits_ / 32;

@@ -162,9 +116,8 @@ void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
    return {x_q, scales_x};
  };
  auto [x_q, scale_x_pre] = quantize(inputs[0], encoder, s);
-  auto [w_q, scale_w_pre] = (inputs[1].dtype() != uint32)
-      ? quantize(inputs[1], encoder, s)
-      : std::make_pair(inputs[1], inputs[2]);
+  auto [w_q, scale_w_pre] = !w_quantized ? quantize(inputs[1], encoder, s)
+                                         : std::make_pair(inputs[1], inputs[2]);

  out.set_data(cu::malloc_async(out.nbytes(), encoder));

@@ -176,8 +129,8 @@ void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
  int K = K_packed * (32 / bits_);

  // Repack scales from linear to tiled layout for tensor cores
-  array scale_x = pad_and_repack_scales(scale_x_pre, encoder, s);
-  array scale_w = pad_and_repack_scales(scale_w_pre, encoder, s);
+  array scale_x = pad_and_swizzle_scales(scale_x_pre, encoder, s);
+  array scale_w = pad_and_swizzle_scales(scale_w_pre, encoder, s);

  bool x_transposed = false;
  bool w_transposed = true; // always transposed
@@ -0,0 +1,50 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/quantized/qqmm_impl.h"
+#include "mlx/backend/cuda/quantized/cublas_qqmm.h"
+
+namespace mlx::core {
+
+void qqmm_impl(
+    cu::CommandEncoder& encoder,
+    int M,
+    int N,
+    int K,
+    bool a_transposed,
+    int64_t lda,
+    bool b_transposed,
+    int64_t ldb,
+    array& out,
+    const array& a,
+    const array& b,
+    const array& a_scale,
+    const array& b_scale,
+    Dtype out_dtype,
+    QuantizationMode mode,
+    float alpha) {
+  // Invoke CublasQQMM
+  std::string qmode = quantization_mode_to_string(mode);
+
+  // Currently only supports non-batched QQMM operations
+  // that covers all use cases for training, we will just collapse (batch,
+  // seq_len) into (tokens)
+  CublasQQMM qqmm(
+      encoder.device(),
+      a_transposed,
+      M,
+      K,
+      lda,
+      b_transposed,
+      K,
+      N,
+      ldb,
+      1, // batch_count
+      0, // a_batch_stride
+      0, // b_batch_stride
+      out_dtype,
+      qmode);
+
+  qqmm.run(encoder, out, a, b, a_scale, b_scale, alpha);
+}
+
+} // namespace mlx::core
@@ -0,0 +1,26 @@
+// Copyright © 2026 Apple Inc.
+#pragma once
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+void qqmm_impl(
+    cu::CommandEncoder& encoder,
+    int M,
+    int N,
+    int K,
+    bool a_transposed,
+    int64_t lda,
+    bool b_transposed,
+    int64_t ldb,
+    array& out,
+    const array& a,
+    const array& b,
+    const array& a_scale,
+    const array& b_scale,
+    Dtype out_dtype,
+    QuantizationMode mode,
+    float alpha = 1.0f);
+
+} // namespace mlx::core
@@ -10,6 +10,11 @@ namespace mlx::core {

 namespace cg = cooperative_groups;

+constexpr int TILE_ROWS = 128;
+constexpr int TILE_COLS = 4;
+constexpr int TILES_PER_LANE = 1;
+constexpr int LANES_PER_BLOCK = 32;
+
 // To pass scales to tensor cores, they need to be repacked into a tiled layout
 // https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout
 // Tiled layout for scale factors is very well described in CUTLASS
@@ -43,118 +48,171 @@ namespace cg = cooperative_groups;
 //          [252, 253, 254, 255],
 //          [380, 381, 382, 383],
 //          [508, 509, 510, 511]]]]],
-__device__ size_t
-scale_tiled_offset(size_t scale_index, size_t num_rows, size_t num_scale_cols) {
-  // Compute the tiled layout offset for scale factors used in tensor cores
-  // This function maps from a linear scale index to the tiled layout expected
-  // by tensor cores (and cublaslt).
-  //
-  // Input: linear scale index (e.g., for a matrix M x K with group_size,
-  //        scale_index ranges from 0 to (M * K/group_size - 1))
-  //
-  // The tiled layout organizes scales into tiles of 128 rows x 4 columns,
-  // where each tile is subdivided into 4 sub-blocks of 32 rows x 4 columns.
-  size_t row = scale_index / num_scale_cols;
-  size_t col = scale_index % num_scale_cols;

-  constexpr size_t rows_per_tile = 128;
-  constexpr size_t rows_per_sub_block = 32;
-  constexpr size_t cols_per_sub_block = 4;
-  constexpr size_t sub_blocks_per_tile = 4; // Vertically stacked
+inline std::tuple<dim3, dim3> get_swizzle_launch_args(
+    size_t M_swizzled,
+    size_t K_swizzled) {
+  constexpr int tiles_per_block = LANES_PER_BLOCK * TILES_PER_LANE;
+  constexpr int warps_per_block = TILE_ROWS / 4; // 128 / 4 = 32

-  // Decompose row position
-  size_t tile_row = row / rows_per_tile; // Which tile row
-  size_t row_in_tile = row % rows_per_tile; // Row within tile
-  size_t sub_block_row =
-      row_in_tile / rows_per_sub_block; // Sub-block within tile
-  size_t row_in_sub_block =
-      row_in_tile % rows_per_sub_block; // Row in sub-block
+  const int num_tiles_k = K_swizzled / TILE_COLS;
+  const int num_tiles_m = M_swizzled / TILE_ROWS;

-  // Decompose column position
-  size_t col_tile = col / cols_per_sub_block; // Which column tile
-  size_t col_in_sub_block = col % cols_per_sub_block; // Column within sub-block
+  dim3 grid;
+  grid.x = cuda::ceil_div(num_tiles_k, tiles_per_block);
+  grid.y = num_tiles_m;
+  grid.z = 1;
+  // Block is always (32, 32) = 1024 threads
+  dim3 block(LANES_PER_BLOCK, warps_per_block, 1);

-  // Compute tile index and offset within tile
-  size_t num_col_tiles = cuda::ceil_div(num_scale_cols, cols_per_sub_block);
-  size_t tile_idx = tile_row * num_col_tiles + col_tile;
-
-  size_t offset_in_tile =
-      (row_in_sub_block * sub_blocks_per_tile * cols_per_sub_block) +
-      (sub_block_row * cols_per_sub_block) + col_in_sub_block;
-
-  constexpr size_t tile_size = rows_per_tile * cols_per_sub_block;
-  return tile_idx * tile_size + offset_in_tile;
+  return std::make_tuple(grid, block);
 }

 namespace cu {

-__global__ void repack_scales(
+__global__ void swizzle_scales(
    const uint8_t* scales_linear,
-    uint8_t* scales_tiled,
-    size_t input_rows,
-    size_t input_cols,
-    size_t output_rows,
-    size_t output_cols) {
+    uint8_t* scales_swizzled,
+    const size_t M,
+    const size_t K,
+    const size_t M_swizzled,
+    const size_t K_swizzled) {
+  constexpr int tile_size = TILE_ROWS * TILE_COLS;
+  constexpr int num_tile_rows_per_thread = 4;
+  constexpr int max_tiles_per_block = LANES_PER_BLOCK * TILES_PER_LANE;
+
+  constexpr int tile_stride = tile_size / 16; // 32 int4s per tile
+
+  // Each thread loads 16 scales from 4 rows (stride 32) and packs them into
+  // int4. For example: thread (0, 0) loads scales at rows 0,32,64,96 of tile 0,
+  // thread (1, 0) loads rows 0,32,64,96 of of tile 1, etc.
+  // The store is strided within a warp (stride 32 int4s), so we first
+  // write to shared memory, then do a coalesced store from shared to global
  auto block_size = cg::this_thread_block().dim_threads();
  auto block_idx = cg::this_thread_block().group_index();
  auto idx_in_block = cg::this_thread_block().thread_index();

-  auto tidx = block_idx.x * block_size.x + idx_in_block.x;
-  auto tidy = block_idx.y * block_size.y + idx_in_block.y;
+  auto tidx = idx_in_block.x;
+  auto tidy = idx_in_block.y;
+  auto linear_tid = tidy * block_size.x + tidx;

-  auto grid_dim_x =
-      cg::this_grid().dim_blocks().x * cg::this_grid().block_index().x;
+  const int bid_x = block_idx.x;
+  const int bid_y = block_idx.y;

-  size_t output_index = tidx + grid_dim_x * size_t(tidy);
-  size_t output_size = output_rows * output_cols;
+  const int K_int = K_swizzled / 4;

-  if (output_index >= output_size) {
-    return;
+  const size_t output_offset = static_cast<size_t>(bid_y) * TILE_ROWS * K_int +
+      static_cast<size_t>(bid_x) * max_tiles_per_block * tile_size / 4;
+  int* output_block = reinterpret_cast<int*>(scales_swizzled) + output_offset;
+
+  const int grid_dim_x = cg::this_grid().dim_blocks().x;
+  const int grid_dim_y = cg::this_grid().dim_blocks().y;
+
+  int remaining = K_int - bid_x * max_tiles_per_block;
+  int tiles_in_block = min(remaining, max_tiles_per_block);
+  bool valid_tile = tidx * TILES_PER_LANE < tiles_in_block;
+
+  __shared__ int4 strided_scales_thread[max_tiles_per_block * tile_stride];
+
+  // Initialize to zero for padding
+  int thread_tile_rows[num_tile_rows_per_thread] = {0};
+
+  if (valid_tile) {
+    const size_t col_base =
+        static_cast<size_t>(bid_x) * max_tiles_per_block * TILE_COLS +
+        tidx * TILE_COLS;
+
+    const bool aligned_k = (K % 4 == 0);
+
+    if (aligned_k) {
+      // fast path: K is aligned, use vectorized loads with stride K/4
+      const int K_stride = K / 4;
+      const size_t block_offset =
+          static_cast<size_t>(bid_y) * TILE_ROWS * K_stride +
+          static_cast<size_t>(bid_x) * max_tiles_per_block;
+      const int* input_block =
+          reinterpret_cast<const int*>(scales_linear) + block_offset;
+// load
+#pragma unroll
+      for (int i = 0; i < num_tile_rows_per_thread; i++) {
+        const size_t row =
+            static_cast<size_t>(bid_y) * TILE_ROWS + i * block_size.x + tidy;
+        const int thread_offset =
+            (i * block_size.x + tidy) * K_stride + tidx * TILES_PER_LANE;
+        if (row < M && col_base + TILE_COLS <= K) {
+          thread_tile_rows[i] = __ldg(input_block + thread_offset);
+        } else if (row < M) {
+// partial tile at K boundary: load byte-by-byte
+#pragma unroll
+          for (int c = 0; c < TILE_COLS; c++) {
+            if (col_base + c < K) {
+              reinterpret_cast<uint8_t*>(&thread_tile_rows[i])[c] =
+                  scales_linear[row * K + col_base + c];
+            }
+          }
+        }
+      }
+    } else {
+#pragma unroll
+      for (int i = 0; i < num_tile_rows_per_thread; i++) {
+        const size_t row =
+            static_cast<size_t>(bid_y) * TILE_ROWS + i * block_size.x + tidy;
+        if (row < M) {
+          const size_t row_start = row * K;
+#pragma unroll
+          for (int c = 0; c < TILE_COLS; c++) {
+            if (col_base + c < K) {
+              reinterpret_cast<uint8_t*>(&thread_tile_rows[i])[c] =
+                  scales_linear[row_start + col_base + c];
+            }
+          }
+        }
+      }
+    }
+    // store to shared with XOR swizzle to avoid bank conflicts
+    int base_idx = tidx * tile_stride + tidy;
+    int xor_bits = (tidy >> 3) & 0x3;
+    int swizzled_idx = base_idx ^ xor_bits;
+    strided_scales_thread[swizzled_idx] =
+        *reinterpret_cast<int4*>(thread_tile_rows);
  }

-  size_t tiled_offset =
-      scale_tiled_offset(output_index, output_rows, output_cols);
+  cg::thread_block block = cg::this_thread_block();
+  cg::sync(block);

-  size_t row = output_index / output_cols;
-  size_t col = output_index % output_cols;
-
-  // Probably this can be done better with 2 separated paths for valid and
-  // padding
-  if (row < input_rows && col < input_cols) {
-    size_t input_index = row * input_cols + col;
-    scales_tiled[tiled_offset] = scales_linear[input_index];
-  } else {
-    // Zero-fill padding region
-    scales_tiled[tiled_offset] = 0;
+  const int total_int4s = tiles_in_block * tile_stride;
+#pragma unroll
+  for (int i = linear_tid; i < total_int4s; i += block_size.x * block_size.y) {
+    int tile_idx = i / tile_stride;
+    int row_idx = i % tile_stride;
+    int base_idx = tile_idx * tile_stride + row_idx;
+    int xor_bits = (row_idx >> 3) & 0x3;
+    int swizzled_idx = base_idx ^ xor_bits;
+    reinterpret_cast<int4*>(output_block)[i] =
+        strided_scales_thread[swizzled_idx];
  }
 }
-
 } // namespace cu

-void repack_scales(
+void swizzle_scales(
    const array& scales,
    array& scales_tiled,
    cu::CommandEncoder& enc,
    const Stream& s) {
  enc.set_input_array(scales);
  enc.set_output_array(scales_tiled);
-
  // Note: scales_tiled is padded to full tiles so if num_rows or num_cols
-  // are not multiples of tile sizes, the extra space is filled with zeros
-
+  // are not multiples of tile sizes
  size_t input_rows = scales.shape(-2);
  size_t input_cols = scales.shape(-1);

  size_t output_rows = scales_tiled.shape(-2);
  size_t output_cols = scales_tiled.shape(-1);
-  size_t output_size = output_rows * output_cols;
-
-  bool large = output_size > UINT_MAX;
-  auto [num_blocks, block_dims] = get_launch_args(
-      output_size, scales_tiled.shape(), scales_tiled.strides(), large);

+  auto [num_blocks, block_dims] =
+      get_swizzle_launch_args(output_rows, output_cols);
  enc.add_kernel_node(
-      cu::repack_scales,
+      cu::swizzle_scales,
      num_blocks,
      block_dims,
      0,
@@ -21,7 +21,7 @@ inline std::pair<int, int> get_padded_scale_dims(int num_rows, int num_cols) {
  return {padded_rows, padded_cols};
 }

-void repack_scales(
+void swizzle_scales(
    const array& scales,
    array& scales_tiled,
    cu::CommandEncoder& enc,
@@ -2,50 +2,48 @@

 #include "mlx/backend/cuda/quantized/quantized.h"
 #include "mlx/backend/cuda/device.h"
-#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/cuda/quantized/qmv.h"
+#include "mlx/backend/cuda/quantized/quantized_utils.h"
 #include "mlx/fast_primitives.h"
+#include "mlx/primitives.h"

 #include <nvtx3/nvtx3.hpp>

 namespace mlx::core {

-namespace {
+void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("QuantizedMatmul::eval_gpu");
+  auto& s = stream();
+  auto& d = cu::device(s.device);
+  auto& enc = d.get_command_encoder(s);

-inline array ensure_row_contiguous(
-    const array& x,
-    cu::CommandEncoder& enc,
-    const Stream& s) {
-  if (!x.flags().row_contiguous) {
-    array x_copy = contiguous_copy_gpu(x, s);
-    enc.add_temporary(x_copy);
-    return x_copy;
-  } else {
-    return x;
+  out.set_data(cu::malloc_async(out.nbytes(), enc));
+
+  // Make sure the last two dims of x and w, s, b are contiguous. This should
+  // be relaxed for x.
+  array x = ensure_row_contiguous_matrix(inputs[0], enc, s);
+  array w = ensure_row_contiguous_matrix(inputs[1], enc, s);
+  array scales = ensure_row_contiguous_matrix(inputs[2], enc, s);
+  std::optional<array> biases = std::nullopt;
+  if (inputs.size() == 4) {
+    biases = ensure_row_contiguous_matrix(inputs[3], enc, s);
+  }
+
+  bool non_batched = w.ndim() == 2 && x.flags().row_contiguous;
+  int K = x.shape(-1);
+  int M = non_batched ? x.size() / K : x.shape(-2);
+  int N = out.shape(-1);
+
+  if (M > 8 || !transpose_ || mode_ == QuantizationMode::Affine) {
+    throw std::runtime_error("QMM NYI");
+  }
+
+  if (transpose_) {
+    fp_qmv(w, scales, x, out, bits_, group_size_, M, N, K, enc);
+    return;
  }
 }

-inline array ensure_row_contiguous_matrix(
-    const array& x,
-    cu::CommandEncoder& enc,
-    const Stream& s) {
-  if (x.ndim() < 2) {
-    if (x.strides()[0] == 1) {
-      return x;
-    }
-  } else {
-    auto stride_0 = x.strides()[x.ndim() - 2];
-    auto stride_1 = x.strides()[x.ndim() - 1];
-    if (stride_0 == x.shape(-1) && stride_1 == 1) {
-      return x;
-    }
-  }
-  array x_copy = contiguous_copy_gpu(x, s);
-  enc.add_temporary(x_copy);
-  return x_copy;
-}
-
-} // namespace
-
 void fast::Quantize::eval_gpu(
    const std::vector<array>& inputs,
    std::vector<array>& outputs) {
@@ -68,7 +66,7 @@ void fast::Quantize::eval_gpu(
      fp_dequantize(wq, scales, w, group_size_, bits_, enc, s);
    }
  } else {
-    auto w = ensure_row_contiguous(inputs[0], enc, s);
+    auto w = ensure_contiguous(inputs[0], enc, s);
    auto& wq = outputs[0];
    auto& scales = outputs[1];

@@ -42,4 +42,12 @@ void fp_dequantize(
    cu::CommandEncoder& enc,
    const Stream& s);

+void fp_quantize_dequantize(
+    const array& w,
+    array& what,
+    int group_size,
+    int bits,
+    cu::CommandEncoder& enc,
+    const Stream& s);
+
 } // namespace mlx::core
@@ -1,9 +1,22 @@
 // Copyright © 2025 Apple Inc.

+#include <cuda_fp4.h>
+#include <cuda_fp8.h>
+
 namespace mlx::core {

 namespace cu {

+inline __device__ float4 dequant_fp8(uint32_t bits) {
+  auto out = *(__nv_fp8x4_e4m3*)(&bits);
+  return out.operator float4();
+}
+
+inline __device__ float4 dequant_fp4(uint16_t bits) {
+  auto out = *(__nv_fp4x4_e2m1*)(&bits);
+  return out.operator float4();
+}
+
 template <int bits, int wsize = 8>
 inline constexpr __device__ short get_pack_factor() {
  return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits);
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
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