remove thrust (#3067 )

Fp qmv (#2984 )
Do not require ConcurrentManagedAccess when not used (#3062 )
2026-01-27 08:54:07 -08:00 · 2026-01-27 06:33:06 -08:00 · 2026-01-27 11:19:20 +09:00 · 2026-01-27 11:01:58 +09:00 · 2026-01-26 16:45:02 -08:00 · 2026-01-27 08:53:26 +09:00
215 changed files with 6277 additions and 2140 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
@@ -4,26 +4,28 @@ description: 'Build and test MLX on macOS'
 runs:
  using: "composite"
  steps:
-    - name: Build
+    - name: Install dependencies
      env:
        DEBUG: 1
        CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
      shell: bash -l {0}
      run: |
+        pip install --upgrade pip
+        pip install cmake setuptools typing_extensions
        pip install -e . -v

    - 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
@@ -75,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%
@@ -18,18 +18,7 @@ runs:
      shell: bash
      run: xcodebuild -showComponent MetalToolchain

-    - name: Install Python
-      shell: sh
-      run: |
-        curl -LsSf https://astral.sh/uv/install.sh | sh
-        $HOME/.local/bin/uv venv --python ${{ inputs.python-version }}
-        source .venv/bin/activate
-        echo PATH=$PATH >> $GITHUB_ENV
-        # Search python packages in .venv
-        echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
-
-    - name: Install build dependencies
-      shell: sh
-      run: |
-        pip install --upgrade pip
-        pip install cmake setuptools typing_extensions
+    - uses: conda-incubator/setup-miniconda@v3
+      with:
+        miniconda-version: "latest"
+        python-version: ${{ inputs.python-version }}
@@ -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,9 +323,6 @@ 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()
@@ -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.                                                 #
@@ -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)


@@ -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
@@ -11,6 +11,7 @@ Transforms
   eval
   async_eval
   compile
+   checkpoint
   custom_function
   disable_compile
   enable_compile
@@ -22,16 +22,56 @@ 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.
+set_target_properties(mlx mlx_version PROPERTIES DEFINE_SYMBOL MLX_EXPORT)
+# Define MLX_STATIC for static libraries.
+if(NOT BUILD_SHARED_LIBS)
+  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)
+  # Disable some MSVC warnings to speed up compilation.
+  target_compile_options(
+    mlx
+    PUBLIC $<$<COMPILE_LANGUAGE:CXX>:/wd4068
+           /wd4244
+           /wd4267
+           /wd4700
+           /wd4804>
+           $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/wd4068
+           -Xcompiler=/wd4244
+           -Xcompiler=/wd4267
+           -Xcompiler=/wd4700
+           -Xcompiler=/wd4804>)
+  # 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;
@@ -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,115 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cpu/device_info.h"
+
+#ifdef __APPLE__
+#include <sys/sysctl.h>
+#include <sys/utsname.h>
+#elif defined(_WIN32)
+#define WIN32_LEAN_AND_MEAN
+#define NOMINMAX
+#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>

@@ -45,12 +47,13 @@ SmallSizePool::SmallSizePool() {

  CHECK_CUDA_ERROR(cudaMallocManaged(&data_, small_pool_size));

-  int device_count = 0;
-  CHECK_CUDA_ERROR(cudaGetDeviceCount(&device_count));
+  int device_count = gpu::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));
+    if (cu::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_;
@@ -294,10 +297,14 @@ 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_;
 }

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

 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 +71,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()) {
@@ -209,6 +215,10 @@ std::pair<int, int> get_graph_limits(Device& d) {
      ops = 50;
      mb = 500;
      break;
+    case 1200: // Consumer Blackwell
+      ops = 100;
+      mb = 1000;
+      break;
    case 1210: // DGX Spark
      ops = 20;
      mb = 25;
@@ -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,23 @@ class Device {
  int compute_capability_minor() const {
    return compute_capability_minor_;
  }
-  cublasLtHandle_t lt_handle() const {
-    return lt_;
-  }
-  cudnnHandle_t cudnn_handle() const {
-    return cudnn_;
+  bool concurrent_managed_access() const {
+    return concurrent_managed_access_ == 1;
  }

 private:
  int device_;
  int compute_capability_major_;
  int compute_capability_minor_;
+  int concurrent_managed_access_;
  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();
@@ -192,7 +194,24 @@ __global__ void event_signal_kernel(AtomicEvent::Atomic* ac, uint64_t value) {
  event_signal(ac, value);
 }

+bool supports_concurrent_managed_access() {
+  static bool concurrent_managed_access = []() {
+    int device_count = gpu::device_count();
+    for (int i = 0; i < device_count; ++i) {
+      if (!cu::device(i).concurrent_managed_access()) {
+        return false;
+      }
+    }
+    return true;
+  }();
+  return concurrent_managed_access;
+}
+
 AtomicEvent::AtomicEvent() {
+  if (!supports_concurrent_managed_access()) {
+    throw std::runtime_error(
+        "Device does not support synchronization in managed memory.");
+  }
  buf_ = std::shared_ptr<Buffer>(
      new Buffer{allocator().malloc(sizeof(Atomic))}, [](Buffer* ptr) {
        allocator().free(*ptr);
@@ -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, bool kEnableTF32>
+struct GemmConfiguration<
+    T,
+    Arch,
+    kAlignmentC,
+    kEnableTF32,
+    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,302 @@
+// 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};
+      uint B = out.size() / (M * N);
+      uint 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()) {
+            auto kernel = fp_qmv_single<T, rows_per_block, n(), 4, 32, true>;
+            if (bits == 8) {
+              kernel = fp_qmv_single<T, rows_per_block, n(), 8, 32, true>;
+            } else if (group_size == 16) {
+              kernel = fp_qmv_single<T, rows_per_block, n(), 4, 16, false>;
+            }
+            encoder.add_kernel_node(
+                kernel,
+                {static_cast<uint>(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(), 4, 32, true>;
+            if (bits == 8) {
+              kernel = fp_qmv_batched<T, rows_per_block, n(), 8, 32, true>;
+            } else if (group_size == 16) {
+              kernel = fp_qmv_batched<T, rows_per_block, n(), 4, 16, false>;
+            }
+            encoder.add_kernel_node(
+                kernel,
+                {static_cast<uint>(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);
@@ -0,0 +1,50 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/gpu/copy.h"
+
+namespace mlx::core {
+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;
+}
+
+inline array
+ensure_contiguous(const array& x, cu::CommandEncoder& enc, const Stream& s) {
+  if (x.flags().row_contiguous || x.flags().col_contiguous) {
+    return x;
+  }
+  array x_copy = contiguous_copy_gpu(x, s);
+  enc.add_temporary(x_copy);
+  return x_copy;
+}
+
+} // namespace mlx::core
@@ -51,9 +51,9 @@ __global__ void rbitsc(
    bool odd,
    uint32_t bytes_per_key) {
  auto grid = cg::this_grid();
-  uint thread_index = grid.thread_rank();
-  uint index_x = thread_index % grid_dims.x;
-  uint index_y = thread_index / grid_dims.x;
+  uint32_t thread_index = grid.thread_rank();
+  uint32_t index_x = thread_index % grid_dims.x;
+  uint32_t index_y = thread_index / grid_dims.x;
  if (index_x >= grid_dims.x || index_y >= grid_dims.y) {
    return;
  }
@@ -94,9 +94,9 @@ __global__ void rbits(
    const __grid_constant__ Shape key_shape,
    const __grid_constant__ Strides key_strides) {
  auto grid = cg::this_grid();
-  uint thread_index = grid.thread_rank();
-  uint index_x = thread_index % grid_dims.x;
-  uint index_y = thread_index / grid_dims.x;
+  uint32_t thread_index = grid.thread_rank();
+  uint32_t index_x = thread_index % grid_dims.x;
+  uint32_t index_y = thread_index / grid_dims.x;
  if (index_x >= grid_dims.x || index_y >= grid_dims.y) {
    return;
  }
@@ -68,8 +68,8 @@ void all_reduce(

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

-  auto get_args = [](size_t size, int N) {
-    int threads = std::min(512UL, (size + N - 1) / N);
+  auto get_args = [](int size, int N) {
+    int threads = std::min(512, (size + N - 1) / N);
    threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
    int reductions_per_step = threads * N;
    size_t steps_needed =
@@ -8,6 +8,7 @@
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
 #include <cub/block/block_load.cuh>
+#include <cub/cub.cuh>

 namespace mlx::core {

@@ -70,14 +70,14 @@ struct Min {
  template <typename T>
  __device__ __forceinline__ T operator()(T a, T b) {
    if constexpr (is_complex_v<T>) {
-      if (isnan(a.real()) || isnan(a.imag())) {
+      if (cuda::std::isnan(a.real()) || cuda::std::isnan(a.imag())) {
        return a;
      }
-      if (isnan(b.real()) || isnan(b.imag())) {
+      if (cuda::std::isnan(b.real()) || cuda::std::isnan(b.imag())) {
        return b;
      }
    } else if constexpr (!cuda::std::is_integral_v<T>) {
-      if (isnan(a) || isnan(b)) {
+      if (cuda::std::isnan(a) || cuda::std::isnan(b)) {
        return cuda::std::numeric_limits<float>::quiet_NaN();
      }
    }
@@ -94,14 +94,14 @@ struct Max {
  template <typename T>
  __device__ __forceinline__ T operator()(T a, T b) {
    if constexpr (is_complex_v<T>) {
-      if (isnan(a.real()) || isnan(a.imag())) {
+      if (cuda::std::isnan(a.real()) || cuda::std::isnan(a.imag())) {
        return a;
      }
-      if (isnan(b.real()) || isnan(b.imag())) {
+      if (cuda::std::isnan(b.real()) || cuda::std::isnan(b.imag())) {
        return b;
      }
    } else if constexpr (!cuda::std::is_integral_v<T>) {
-      if (isnan(a) || isnan(b)) {
+      if (cuda::std::isnan(a) || cuda::std::isnan(b)) {
        return cuda::std::numeric_limits<float>::quiet_NaN();
      }
    }
@@ -30,7 +30,7 @@ __device__ void rope_single_impl(
  float sintheta = sin(theta);

  // Compute the input and output indices
-  uint index_1, index_2;
+  uint32_t index_1, index_2;
  if (traditional) {
    index_1 = 2 * pos.x + pos.y * stride;
    index_2 = index_1 + 1;
@@ -269,7 +269,7 @@ void sdpa_cudnn(
    bool output_logsumexp,
    Stream s) {
  auto& encoder = cu::get_command_encoder(s);
-  auto handle = encoder.device().cudnn_handle();
+  auto handle = encoder.device().get_cudnn_handle();

  malloc_with_same_layout(encoder, o, q);

@@ -327,7 +327,7 @@ void sdpa_backward_cudnn(
    array& d_v,
    Stream s) {
  auto& encoder = cu::get_command_encoder(s);
-  auto handle = encoder.device().cudnn_handle();
+  auto handle = encoder.device().get_cudnn_handle();

  malloc_with_same_layout(encoder, d_q, q);
  malloc_with_same_layout(encoder, d_k, k);
@@ -1,5 +1,8 @@
 // Copyright © 2025 Apple Inc.

+// Required for using M_LOG2E in MSVC.
+#define _USE_MATH_DEFINES
+
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/device/config.h"
 #include "mlx/backend/cuda/device/utils.cuh"
@@ -227,21 +227,21 @@ __global__ void strided_scan(
  // Compute offsets.
  int64_t offset = (grid.block_rank() / stride_blocks) * axis_size * stride;
  int64_t global_index_x = (grid.block_rank() % stride_blocks) * BN;
-  uint read_offset_y = (block.thread_rank() * N_READS) / BN;
-  uint read_offset_x = (block.thread_rank() * N_READS) % BN;
-  uint scan_offset_y = warp.thread_rank();
-  uint scan_offset_x = warp.meta_group_rank() * n_scans;
+  uint32_t read_offset_y = (block.thread_rank() * N_READS) / BN;
+  uint32_t read_offset_x = (block.thread_rank() * N_READS) % BN;
+  uint32_t scan_offset_y = warp.thread_rank();
+  uint32_t scan_offset_x = warp.meta_group_rank() * n_scans;

-  uint stride_limit = stride - global_index_x;
+  uint32_t stride_limit = stride - global_index_x;
  in += offset + global_index_x + read_offset_x;
  out += offset + global_index_x + read_offset_x;
  U* read_into = read_buffer + read_offset_y * BN_pad + read_offset_x;
  U* read_from = read_buffer + scan_offset_y * BN_pad + scan_offset_x;

-  for (uint j = 0; j < axis_size; j += BM) {
+  for (uint32_t j = 0; j < axis_size; j += BM) {
    // Calculate the indices for the current thread.
-    uint index_y = j + read_offset_y;
-    uint check_index_y = index_y;
+    uint32_t index_y = j + read_offset_y;
+    uint32_t check_index_y = index_y;
    if (reverse) {
      index_y = axis_size - 1 - index_y;
    }
@@ -395,7 +395,7 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
    using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
    dispatch_scan_ops(reduce_type_, [&](auto scan_op_tag) {
      using Op = MLX_GET_TYPE(scan_op_tag);
-      if constexpr (supports_scan_op<Op, T>) {
+      if constexpr (supports_scan_op<Op, T>()) {
        using U = typename cu::ScanResult<Op, T>::type;
        dispatch_bool(inclusive_, [&](auto inclusive) {
          dispatch_bool(reverse_, [&](auto reverse) {
@@ -454,11 +454,12 @@ void Scan::eval_gpu(const std::vector<array>& inputs, array& out) {
          });
        });
      } else {
-        throw std::runtime_error(fmt::format(
-            "Can not do scan op {} on inputs of {} with result of {}.",
-            op_to_string<Op>(),
-            dtype_to_string(in.dtype()),
-            dtype_to_string(out.dtype())));
+        throw std::runtime_error(
+            fmt::format(
+                "Can not do scan op {} on inputs of {} with result of {}.",
+                op_to_string<Op>(),
+                dtype_to_string(in.dtype()),
+                dtype_to_string(out.dtype())));
      }
    });
  });
@@ -191,11 +191,12 @@ void unary_op_gpu_inplace(
          }
        });
      } else {
-        throw std::runtime_error(fmt::format(
-            "Can not do unary op {} on input of {} with output of {}.",
-            op,
-            dtype_to_string(in.dtype()),
-            dtype_to_string(out.dtype())));
+        throw std::runtime_error(
+            fmt::format(
+                "Can not do unary op {} on input of {} with output of {}.",
+                op,
+                dtype_to_string(in.dtype()),
+                dtype_to_string(out.dtype())));
      }
    });
  });
@@ -5,6 +5,7 @@
 #include "mlx/dtype_utils.h"

 #include <fmt/format.h>
+#include <cuda/cmath>
 #include <vector>

 namespace mlx::core {
@@ -11,7 +11,7 @@
 namespace mlx::core {

 template <typename T>
-inline uint max_occupancy_block_dim(T kernel) {
+inline uint32_t max_occupancy_block_dim(T kernel) {
  int _, block_dim;
  if constexpr (std::is_same_v<T, CUfunction>) {
    CHECK_CUDA_ERROR(
@@ -1,9 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-namespace mlx::core::gpu {
-
-bool is_available();
-
-} // namespace mlx::core::gpu
@@ -0,0 +1,36 @@
+// Copyright © 2026 Apple Inc.
+
+#pragma once
+
+#include <string>
+#include <unordered_map>
+#include <variant>
+
+#include "mlx/api.h"
+
+namespace mlx::core::gpu {
+
+MLX_API bool is_available();
+
+/**
+ * Get the number of available GPU devices.
+ */
+MLX_API int device_count();
+
+/**
+ * Get information about a GPU device.
+ *
+ * Returns a map of device properties. Keys vary by backend:
+ *   - device_name (string): Device name
+ *   - architecture (string): Architecture identifier
+ *   - total_memory/memory_size (size_t): Total device memory
+ *   - free_memory (size_t): Available memory (CUDA only)
+ *   - uuid (string): Device UUID (CUDA only)
+ *   - pci_bus_id (string): PCI bus ID (CUDA only)
+ *   - compute_capability_major/minor (size_t): Compute capability (CUDA only)
+ */
+MLX_API const
+    std::unordered_map<std::string, std::variant<std::string, size_t>>&
+    device_info(int device_index = 0);
+
+} // namespace mlx::core::gpu
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
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