Use uv in macOS CI

[CUDA] Fix half type matmul in cutlass kernels (#3469 )
Clearer error when shape dimension overflows int32 (#3425 )
2026-05-06 15:41:43 +09:00 · 2026-05-06 08:35:53 +09:00 · 2026-05-05 09:53:36 +09:00 · 2026-05-05 09:22:23 +09:00 · 2026-05-05 09:06:23 +09:00 · 2026-04-28 09:39:56 -07:00
356 changed files with 21847 additions and 6442 deletions
@@ -18,7 +18,7 @@ runs:
      env:
        CMAKE_ARGS: -DMLX_BUILD_CUDA=ON
      run: |
-        pip install auditwheel build patchelf setuptools
+        pip install auditwheel "build<=1.4.2" patchelf setuptools
        python setup.py clean --all
        MLX_BUILD_STAGE=2 python -m build -w

@@ -25,7 +25,7 @@ runs:
    - name: Build Python package
      shell: bash
      run: |
-        pip install auditwheel patchelf build
+        pip install auditwheel patchelf "build<=1.4.2"
        python setup.py clean --all
        MLX_BUILD_STAGE=1 python -m build -w
        auditwheel repair dist/mlx-*.whl \
@@ -21,7 +21,7 @@ runs:
      run: |
        if ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }} ; then
          # There is no GPU in arm64 runner, use a common arch.
-          CMAKE_ARGS="$CMAKE_ARGS -DMLX_CUDA_ARCHITECTURES=90a"
+          CMAKE_ARGS="$CMAKE_ARGS -DMLX_CUDA_ARCHITECTURES=80"
          # Can not build tests and stubs when the built executables can not run.
          CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF -DMLX_BUILD_PYTHON_STUBS=OFF"
        fi
@@ -21,7 +21,7 @@ runs:
        DEVELOPER_DIR: /Applications/Xcode-latest.app
        MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }}
      run: |
-        pip install build
+        uv pip install build
        python setup.py clean --all
        MLX_BUILD_STAGE=1 python -m build -w

@@ -4,59 +4,72 @@ description: 'Build and test MLX on macOS'
 runs:
  using: "composite"
  steps:
-    - name: Install dependencies
+    - name: Install Python package
      env:
        DEBUG: 1
        CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
-      shell: bash -l {0}
+      shell: bash
      run: |
-        pip install --upgrade pip
-        pip install cmake setuptools typing_extensions
-        pip install -e . -v
+        echo "::group::Install Python package"
+        uv pip install -e ".[dev]" -v
+        echo "::endgroup::"

    - name: Install tests dependencies
-      shell: bash -l {0}
+      shell: bash
      run: |
-        pip install numpy torch tensorflow
+        echo "::group::Install tests dependencies"
+        uv pip install tensorflow
+        echo "::endgroup::"

    - name: Run Python tests
-      shell: bash -l {0}
+      shell: bash
      env:
        LOW_MEMORY: 1
      run: |
+        echo "::group::Run Python tests"
        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
-    
+        echo "::endgroup::"
+
    - name: Build example extension
-      shell: bash -l {0}
+      shell: bash
      run: |
+        echo "::group::Build example extension"
        cd examples/extensions
-        pip install -r requirements.txt
-        python setup.py build_ext --inplace
-        python test.py
-    
+        uv pip install -r requirements.txt
+        uv run --no-project setup.py build_ext --inplace
+        uv run --no-project test.py
+        echo "::endgroup::"
+
    - name: Build CPP only
-      shell: bash -l {0}
+      shell: bash
      run: |
+        echo "::group::Build CPP only"
        mkdir -p build
        cd build
        cmake ..
        make -j $(sysctl -n hw.ncpu)
-    
+        echo "::endgroup::"
+
    - name: Run CPP tests
-      shell: bash -l {0}
+      shell: bash
      env:
        DEVICE: gpu
        METAL_DEVICE_WRAPPER_TYPE: 1
        METAL_DEBUG_ERROR_MODE: 0
-      run: ./build/tests/tests
-    
-    - name: Build small binary with JIT
-      shell: bash -l {0}
      run: |
+        echo "::group::Run CPP tests"
+        ./build/tests/tests
+        ./build/tests/test_teardown
+        echo "::endgroup::"
+
+    - name: Build small binary with JIT
+      shell: bash
+      run: |
+        echo "::group::Build small binary with JIT"
        mkdir -p build
        cd build
        cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \
@@ -66,15 +79,18 @@ runs:
          -DMLX_BUILD_GGUF=OFF \
          -DMLX_METAL_JIT=ON
        make -j $(sysctl -n hw.ncpu)
-    
+        echo "::endgroup::"
+
    - name: Run Python tests with JIT
-      shell: bash -l {0}
+      shell: bash
      env:
        LOW_MEMORY: 1
        DEVICE: gpu
        METAL_DEVICE_WRAPPER_TYPE: 1
        METAL_DEBUG_ERROR_MODE: 0
      run: |
+        echo "::group::Run Python tests with JIT"
        CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
-          pip install -e . -v
+          uv pip install -e . -v
        python -m unittest discover -v python/tests
+        echo "::endgroup::"
@@ -14,6 +14,9 @@ inputs:
    description: 'Whether to enable ccache'
    required: false
    default: 'true'
+  ccache-key:
+    required: false
+    default: 'ccache'

 runs:
  using: "composite"
@@ -33,7 +36,7 @@ runs:
      if: ${{ inputs.use-ccache == 'true' }}
      uses: hendrikmuhs/ccache-action@v1.2
      with:
-        key: ccache-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}
+        key: ${{ inputs.ccache-key }}-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}
        max-size: 1GB
        # ccache-action bug: running "apt-get update" fails on large arm runner.
        update-package-index: false
@@ -54,6 +57,12 @@ runs:
        echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
        echo "::endgroup::"

+    - name: Set swap space
+      if: ${{ startsWith(inputs.toolkit, 'cuda') }}
+      uses: pierotofy/set-swap-space@fc79b3f67fa8a838184ce84a674ca12238d2c761
+      with:
+        swap-size-gb: 16
+
    - name: Install CUDA toolkit
      if: ${{ startsWith(inputs.toolkit, 'cuda') }}
      shell: bash
@@ -13,12 +13,20 @@ runs:
    - name: Install Homebrew packages
      shell: sh
      run: /opt/homebrew/bin/brew install openmpi
-    
+
    - name: Verify MetalToolchain installed
      shell: bash
      run: xcodebuild -showComponent MetalToolchain

-    - uses: conda-incubator/setup-miniconda@v3
-      with:
-        miniconda-version: "latest"
-        python-version: ${{ inputs.python-version }}
+    - uses: astral-sh/setup-uv@v7
+
+    - name: Setup Python venv
+      shell: bash
+      run: |
+        echo "::group::Setup Python venv"
+        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
+        echo "::endgroup::"
@@ -65,5 +65,5 @@ runs:
        DEVICE: gpu
      run: |
        echo "::group::CPP tests - GPU"
-        ./build/tests/tests -sfe="*fft_tests.cpp,*linalg_tests.cpp"
+        ./build/tests/tests -sfe="*linalg_tests.cpp"
        echo "::endgroup::"
@@ -17,4 +17,5 @@ runs:
      run: |
        echo "::group::CPP tests - CPU"
        ./build/tests.exe -tce="*gguf*,test random uniform"
+        ./build/test_teardown.exe
        echo "::endgroup::"
@@ -25,4 +25,4 @@ jobs:
    steps:
      - name: Deploy to GitHub Pages
        id: deployment
-        uses: actions/deploy-pages@v4
+        uses: actions/deploy-pages@v5
@@ -23,14 +23,14 @@ jobs:
          build-backend: ${{ matrix.python-version == '3.10' }}
          arch: "x86_64"
      - name: Upload mlx artifacts
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
        with:
          name: linux-wheels-${{ matrix.python_version }}
          path: wheelhouse/mlx-*.whl
          retention-days: 7
      - name: Upload mlx-cpu artifacts
        if: matrix.python_version == '3.10'
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
        with:
          name: mlx-cpu
          path: wheelhouse/mlx_cpu-*.whl
@@ -85,20 +85,24 @@ jobs:

  build_cuda_release:
    if: github.repository == 'ml-explore/mlx'
-    runs-on: ubuntu-22-large
+    strategy:
+      matrix:
+        arch: ['x86_64', 'aarch64']
+        toolkit: ['cuda-12.9', 'cuda-13.0']
+    runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22-large' || 'ubuntu-22-large-arm' }}
    steps:
      - uses: actions/checkout@v6
      - uses: ./.github/actions/setup-linux
        with:
-          toolkit: 'cuda-12.9'
+          toolkit: ${{ matrix.toolkit }}
+          ccache-key: 'ccache-release'
      - name: Build Python package
        uses: ./.github/actions/build-cuda-release
        with:
-          toolkit: 'cuda-12.9'
-          arch: 'x86_64'
+          arch: ${{ matrix.arch }}
      - name: Upload artifacts
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
        with:
-          name: mlx-cuda
+          name: mlx-${{ matrix.toolkit }}-${{ matrix.arch }}
          path: wheelhouse/mlx_cuda_*.whl
          retention-days: 7
@@ -41,7 +41,7 @@ jobs:
    steps:
      - name: Deploy to GitHub Pages
        id: deployment
-        uses: actions/deploy-pages@v4
+        uses: actions/deploy-pages@v5

  build_linux_release:
    if: github.repository == 'ml-explore/mlx'
@@ -64,7 +64,7 @@ jobs:
          build-backend: ${{ matrix.python_version == '3.10' }}
          arch: ${{ matrix.arch }}
      - name: Upload MLX artifacts
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
        with:
          overwrite: true
          name: linux-wheels-${{ matrix.python_version }}-${{ matrix.arch }}
@@ -72,7 +72,7 @@ jobs:
          if-no-files-found: error
      - name: Upload CPU artifacts
        if: matrix.python_version == '3.10'
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
        with:
          overwrite: true
          name: mlx-cpu-${{ matrix.arch }}
@@ -93,13 +93,8 @@ jobs:
      - uses: ./.github/actions/setup-macos
        with:
          python-version: ${{ matrix.python-version }}
-
-      - name: Install dependencies
-        shell: bash -l {0}
-        run: |
-          pip install --upgrade pip
-          pip install cmake setuptools typing_extensions
-          pip install -e . -v
+      - name: Install Python package
+        run: uv pip install -e . -v
      - name: Build macOS 14 package
        uses: ./.github/actions/build-macos-release
        with:
@@ -116,7 +111,7 @@ jobs:
          macos-target: 26.0
          build-backend: ${{ matrix.python-version == '3.10' }}
      - name: Upload MLX artifacts
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
        with:
          overwrite: true
          name: mac-wheels-${{ matrix.python-version }}
@@ -124,7 +119,7 @@ jobs:
          if-no-files-found: error
      - name: Upload Metal artifacts
        if: matrix.python-version == '3.10'
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
        with:
          overwrite: true
          name: mlx-metal
@@ -146,13 +141,13 @@ jobs:
      - uses: ./.github/actions/setup-linux
        with:
          toolkit: ${{ matrix.toolkit }}
-          use-ccache: false
+          ccache-key: 'ccache-release'
      - name: Build Python package
        uses: ./.github/actions/build-cuda-release
        with:
          arch: ${{ matrix.arch }}
      - name: Upload artifacts
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
        with:
          overwrite: true
          name: mlx-${{ matrix.toolkit }}-${{ matrix.arch }}
@@ -169,12 +164,12 @@ jobs:
      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
      url: https://pypi.org/p/mlx
    steps:
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
        with:
          pattern: linux-wheels-*
          merge-multiple: true
          path: dist
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
        with:
          pattern: mac-wheels-*
          merge-multiple: true
@@ -197,7 +192,7 @@ jobs:
      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
      url: https://pypi.org/p/mlx-cuda
    steps:
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
        with:
          pattern: mlx-cuda-*
          merge-multiple: true
@@ -220,7 +215,7 @@ jobs:
      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
      url: https://pypi.org/p/mlx-cpu
    steps:
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
        with:
          pattern: mlx-cpu-*
          merge-multiple: true
@@ -243,7 +238,7 @@ jobs:
      name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }}
      url: https://pypi.org/p/mlx-metal
    steps:
-      - uses: actions/download-artifact@v7
+      - uses: actions/download-artifact@v8
        with:
          name: mlx-metal
          path: dist
@@ -156,6 +156,10 @@ if(MLX_BUILD_CUDA)
  enable_language(CUDA)
  find_package(CUDAToolkit REQUIRED)
  find_package(CUDNN REQUIRED)
+  if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "13.1" AND CUDAToolkit_VERSION
+                                                          VERSION_LESS "13.2")
+    message(FATAL_ERROR "CUDA Toolkit 13.1 is not supported.")
+  endif()
 endif()

 if(MLX_BUILD_METAL)
@@ -317,6 +321,15 @@ FetchContent_MakeAvailable(json)
 target_include_directories(
  mlx PRIVATE $<BUILD_INTERFACE:${json_SOURCE_DIR}/single_include/nlohmann>)

+# Add standalone JACCL library (RDMA over Thunderbolt distributed backend)
+if(MLX_BUILD_CPU
+   AND ${CMAKE_SYSTEM_NAME} MATCHES "Darwin"
+   AND DEFINED MACOS_SDK_VERSION
+   AND MACOS_SDK_VERSION GREATER_EQUAL 26.2)
+  add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx/distributed/jaccl/lib
+                   ${CMAKE_BINARY_DIR}/jaccl)
+endif()
+
 add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx)

 target_include_directories(
@@ -329,7 +342,7 @@ else()
  FetchContent_Declare(
    fmt
    GIT_REPOSITORY https://github.com/fmtlib/fmt.git
-    GIT_TAG 10.2.1
+    GIT_TAG 12.1.0
    EXCLUDE_FROM_ALL)
  FetchContent_MakeAvailable(fmt)
 endif()
@@ -344,7 +357,7 @@ if(MLX_BUILD_PYTHON_BINDINGS)
  FetchContent_Declare(
    nanobind
    GIT_REPOSITORY https://github.com/wjakob/nanobind.git
-    GIT_TAG v2.10.2
+    GIT_TAG v2.12.0
    GIT_SHALLOW TRUE
    EXCLUDE_FROM_ALL)
  FetchContent_MakeAvailable(nanobind)
@@ -0,0 +1,193 @@
+# Copyright © 2025 Apple Inc.
+
+import argparse
+import time
+
+import mlx.core as mx
+import numpy as np
+
+MLX_DTYPES = {
+    "float16": mx.float16,
+    "bfloat16": mx.bfloat16,
+    "float32": mx.float32,
+}
+
+
+def parse_cases(cases):
+    parsed = []
+    for spec in cases.split(","):
+        parts = spec.split("x")
+        m, n, k, bs = int(parts[0]), int(parts[1]), int(parts[2]), int(parts[3])
+        sparsity = float(parts[4]) if len(parts) > 4 else 0.5
+        parsed.append((m, n, k, bs, sparsity))
+    return parsed
+
+
+def make_masks(m, n, k, block_size, sparsity, rng):
+    """Create block masks with given sparsity (fraction of blocks zeroed)."""
+    tm = (m + block_size - 1) // block_size
+    tn = (n + block_size - 1) // block_size
+    tk = (k + block_size - 1) // block_size
+
+    lhs_mask = (rng.random((tm, tk)) >= sparsity).astype(np.bool_)
+    rhs_mask = (rng.random((tk, tn)) >= sparsity).astype(np.bool_)
+    out_mask = (rng.random((tm, tn)) >= sparsity).astype(np.bool_)
+    return lhs_mask, rhs_mask, out_mask
+
+
+def mlx_naive_block_masked_mm(a, b, block_size, out_mask, lhs_mask, rhs_mask):
+    """MLX naive: expand masks and use regular matmul."""
+    M, K = a.shape[-2], a.shape[-1]
+    N = b.shape[-1]
+
+    def expand(mask, rows, cols):
+        e = mx.repeat(mx.repeat(mask, block_size, axis=-2), block_size, axis=-1)
+        return e[..., :rows, :cols]
+
+    a_masked = a * expand(lhs_mask, M, K)
+    b_masked = b * expand(rhs_mask, K, N)
+    c = a_masked @ b_masked
+    c = c * expand(out_mask, M, N)
+    return c
+
+
+def bench_mlx(fn, warmup, iters):
+    for _ in range(warmup):
+        y = fn()
+        mx.eval(y)
+    mx.synchronize()
+
+    start = time.perf_counter()
+    for _ in range(iters):
+        y = fn()
+        mx.eval(y)
+    mx.synchronize()
+    return (time.perf_counter() - start) * 1e3 / iters
+
+
+def print_table(headers, rows):
+    widths = [len(h) for h in headers]
+    for row in rows:
+        for i, cell in enumerate(row):
+            widths[i] = max(widths[i], len(cell))
+
+    def fmt_row(row):
+        return (
+            "| "
+            + " | ".join(f"{cell:<{widths[i]}}" for i, cell in enumerate(row))
+            + " |"
+        )
+
+    sep = "|-" + "-|-".join("-" * w for w in widths) + "-|"
+    print(fmt_row(headers))
+    print(sep)
+    for row in rows:
+        print(fmt_row(row))
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Benchmark block_masked_mm vs naive expand+matmul"
+    )
+    parser.add_argument(
+        "--cases",
+        default=(
+            "256x256x256x32x0.5,"
+            "512x512x512x32x0.5,"
+            "1024x1024x1024x32x0.5,"
+            "1024x1024x1024x64x0.5,"
+            "2048x2048x2048x64x0.5,"
+            "256x256x256x32x0.0,"
+            "1024x1024x1024x32x0.0,"
+            "1024x1024x1024x32x0.9"
+        ),
+        help="Comma-separated MxNxKxBSxSparsity list. Sparsity=fraction of blocks zeroed.",
+    )
+    parser.add_argument(
+        "--dtype",
+        default="float32",
+        choices=["float16", "bfloat16", "float32"],
+    )
+    parser.add_argument("--warmup", type=int, default=10)
+    parser.add_argument("--iters", type=int, default=50)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--no-check", action="store_true")
+    args = parser.parse_args()
+
+    mlx_dtype = MLX_DTYPES[args.dtype]
+
+    print(f"dtype={args.dtype}  warmup={args.warmup}  iters={args.iters}")
+
+    headers = [
+        "Case (MxNxKxBS)",
+        "Sparsity",
+        "MLX ms",
+        "Naive ms",
+        "Speedup",
+    ]
+    if not args.no_check:
+        headers.append("Max err")
+    rows = []
+
+    cases = parse_cases(args.cases)
+    for idx, (m, n, k, bs, sparsity) in enumerate(cases):
+        rng = np.random.default_rng(args.seed + idx)
+        a_np = rng.standard_normal((m, k)).astype(np.float32)
+        b_np = rng.standard_normal((k, n)).astype(np.float32)
+        lhs_mask_np, rhs_mask_np, out_mask_np = make_masks(m, n, k, bs, sparsity, rng)
+
+        a_mx = mx.array(a_np, dtype=mlx_dtype)
+        b_mx = mx.array(b_np, dtype=mlx_dtype)
+        lhs_mask_mx = mx.array(lhs_mask_np)
+        rhs_mask_mx = mx.array(rhs_mask_np)
+        out_mask_mx = mx.array(out_mask_np)
+        mx.eval(a_mx, b_mx, lhs_mask_mx, rhs_mask_mx, out_mask_mx)
+
+        # Correctness check: block_masked_mm vs naive expand+matmul
+        err_str = ""
+        if not args.no_check:
+            y_op = mx.block_masked_mm(
+                a_mx, b_mx, bs, out_mask_mx, lhs_mask_mx, rhs_mask_mx
+            )
+            y_naive = mlx_naive_block_masked_mm(
+                a_mx, b_mx, bs, out_mask_mx, lhs_mask_mx, rhs_mask_mx
+            )
+            mx.eval(y_op, y_naive)
+            err = float(mx.max(mx.abs(y_op - y_naive)).item())
+            err_str = f"{err:.2e}"
+
+        # Benchmark
+        t_mlx = bench_mlx(
+            lambda: mx.block_masked_mm(
+                a_mx, b_mx, bs, out_mask_mx, lhs_mask_mx, rhs_mask_mx
+            ),
+            args.warmup,
+            args.iters,
+        )
+        t_naive = bench_mlx(
+            lambda: mlx_naive_block_masked_mm(
+                a_mx, b_mx, bs, out_mask_mx, lhs_mask_mx, rhs_mask_mx
+            ),
+            args.warmup,
+            args.iters,
+        )
+        speedup = f"{t_naive / t_mlx:.2f}x" if t_mlx > 0 else "-"
+
+        row = [
+            f"{m}x{n}x{k}x{bs}",
+            f"{sparsity:.0%}",
+            f"{t_mlx:.3f}",
+            f"{t_naive:.3f}",
+            speedup,
+        ]
+        if not args.no_check:
+            row.append(err_str)
+        rows.append(row)
+
+    print_table(headers, rows)
+    if not args.no_check:
+        print("err: max|block_masked_mm - naive_expand_matmul|")
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,152 @@
+import math
+import time
+
+import mlx.core as mx
+import numpy as np
+import torch
+
+N_warmup = 2
+N_iter_bench = 10
+N_iter_func = 10
+
+
+def bench(f, a, b, b_prime):
+    for i in range(N_warmup):
+        f(a, b, b_prime)
+    torch.mps.synchronize()
+
+    s = time.perf_counter_ns()
+    for i in range(N_iter_bench):
+        f(a, b, b_prime)
+    e = time.perf_counter_ns()
+    return (e - s) * 1e-9
+
+
+def make_mx_conv_3D(strides=(1, 1, 1), padding=(0, 0, 0), groups=1):
+    def mx_conv_3D(a, b, b_prime):
+        y = a
+        for i in range(N_iter_func):
+            y = mx.conv3d(y, b, stride=strides, padding=padding, groups=groups)
+            y = mx.conv3d(y, b_prime, stride=strides, padding=padding, groups=groups)
+        mx.eval(y)
+        return y
+
+    return mx_conv_3D
+
+
+def make_pt_conv_3D(strides=(1, 1, 1), padding=(0, 0, 0), groups=1):
+    @torch.no_grad()
+    def pt_conv_3D(a, b, b_prime):
+        y = a
+        for i in range(N_iter_func):
+            y = torch.conv3d(y, b, stride=strides, padding=padding, groups=groups)
+            y = torch.conv3d(y, b_prime, stride=strides, padding=padding, groups=groups)
+        torch.mps.synchronize()
+        return y
+
+    return pt_conv_3D
+
+
+def bench_shape(N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype):
+    scale = 1.0 / math.sqrt(kD * kH * kW * C)
+    a_np = np.random.uniform(0, 0.5, (N, D, H, W, C))
+    b_np = np.random.uniform(-scale, scale, (O, kD, kH, kW, int(C / groups)))
+    b_prime_np = np.random.uniform(-scale, scale, (C, kD, kH, kW, int(O / groups)))
+
+    a_np, b_np, b_prime_np = map(lambda x: x.astype(np_dtype), (a_np, b_np, b_prime_np))
+    a_mx, b_mx, b_prime_mx = map(lambda x: mx.array(x), (a_np, b_np, b_prime_np))
+    a_pt, b_pt, b_prime_pt = map(
+        lambda x: torch.from_numpy(x.transpose(0, 4, 1, 2, 3)).to("mps"),
+        (a_np, b_np, b_prime_np),
+    )
+
+    torch.mps.synchronize()
+
+    f_mx = make_mx_conv_3D(strides, padding, groups)
+    f_pt = make_pt_conv_3D(strides, padding, groups)
+
+    time_torch = bench(f_pt, a_pt, b_pt, b_prime_pt)
+    time_mlx = bench(f_mx, a_mx, b_mx, b_prime_mx)
+
+    # Measure MLX memory
+    mx.clear_cache()
+    mx.reset_peak_memory()
+    y = mx.conv3d(a_mx, b_mx, stride=strides, padding=padding, groups=groups)
+    mx.eval(y)
+    mlx_peak_mb = mx.get_peak_memory() / 1024**2
+    mlx_active_mb = mx.get_active_memory() / 1024**2
+    del y
+
+    # Measure PyTorch MPS memory
+    torch.mps.synchronize()
+    torch.mps.empty_cache()
+    y = torch.conv3d(a_pt, b_pt, stride=strides, padding=padding, groups=groups)
+    torch.mps.synchronize()
+    pt_current_mb = torch.mps.current_allocated_memory() / 1024**2
+    pt_driver_mb = torch.mps.driver_allocated_memory() / 1024**2
+    del y
+
+    out_mx = mx.conv3d(a_mx, b_mx, stride=strides, padding=padding, groups=groups)
+    out_pt = torch.conv3d(
+        a_pt.to("cpu"), b_pt.to("cpu"), stride=strides, padding=padding, groups=groups
+    )
+    out_pt = torch.permute(out_pt, (0, 2, 3, 4, 1))
+    out_pt = out_pt.numpy(force=True)
+
+    atol = 2e-5 if np_dtype == np.float32 else 5e-4
+
+    if not np.allclose(out_pt, out_mx, atol=atol):
+        print(
+            f"Failed at {(N, D, H, W, C)}, {(O, kD, kH, kW, C)} "
+            f"[strides = {strides}, padding = {padding}, groups = {groups}] "
+            f"with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}"
+        )
+
+    return time_mlx, time_torch, mlx_peak_mb, mlx_active_mb, pt_current_mb, pt_driver_mb
+
+
+if __name__ == "__main__":
+    dtypes = ("float16", "float32")
+    shapes = (
+        # (C % 16 == 0)
+        (4, 16, 16, 16, 32, 3, 3, 3, 32, (1, 1, 1), (1, 1, 1), 1),
+        (4, 16, 16, 16, 64, 3, 3, 3, 64, (1, 1, 1), (1, 1, 1), 1),
+        (4, 16, 16, 16, 128, 3, 3, 3, 128, (1, 1, 1), (1, 1, 1), 1),
+        (4, 32, 32, 32, 64, 3, 3, 3, 64, (1, 1, 1), (1, 1, 1), 1),
+        (4, 32, 32, 32, 128, 3, 3, 3, 128, (1, 1, 1), (1, 1, 1), 1),
+        # Larger spatial dims
+        (2, 64, 64, 64, 32, 3, 3, 3, 64, (1, 1, 1), (1, 1, 1), 1),
+        (1, 64, 64, 64, 64, 3, 3, 3, 128, (1, 1, 1), (1, 1, 1), 1),
+        # Strided
+        (4, 32, 32, 32, 64, 3, 3, 3, 128, (2, 2, 2), (1, 1, 1), 1),
+        # Asymmetric kernels
+        (4, 32, 32, 32, 64, 3, 1, 1, 128, (1, 1, 1), (1, 0, 0), 1),
+        (4, 32, 32, 32, 64, 1, 3, 3, 128, (1, 1, 1), (0, 1, 1), 1),
+        # (C % 16 != 0)
+        (4, 16, 16, 16, 21, 3, 3, 3, 21, (1, 1, 1), (1, 1, 1), 1),
+        (4, 16, 16, 16, 55, 3, 3, 3, 55, (1, 1, 1), (1, 1, 1), 1),
+        (4, 32, 32, 32, 55, 3, 3, 3, 55, (1, 1, 1), (1, 1, 1), 1),
+        (4, 16, 16, 16, 3, 3, 3, 3, 32, (1, 1, 1), (1, 1, 1), 1),
+    )
+
+    for dtype in dtypes:
+        print(f"\n{'=' * 120}" f"\n  dtype: {dtype}" f"\n{'=' * 120}")
+        print(
+            f"{'(N,   D,   H,   W,   C)':<26s} {'(  O, kD, kH, kW,   C)':<24s} "
+            f"{'stride':<12s} {'pads':<12s} {'groups':>6s} "
+            f"{'diff%':>7s}  "
+            f"{'MLX peak':>9s} {'MLX act':>8s} {'PT cur':>8s} {'PT drv':>8s}"
+        )
+        for N, D, H, W, C, kD, kH, kW, O, strides, padding, groups in shapes:
+            np_dtype = getattr(np, dtype)
+            time_mlx, time_torch, mlx_peak, mlx_act, pt_cur, pt_drv = bench_shape(
+                N, D, H, W, C, kD, kH, kW, O, strides, padding, groups, np_dtype
+            )
+            diff = time_torch / time_mlx - 1.0
+
+            print(
+                f"({N}, {D:3d}, {H:3d}, {W:3d}, {C:3d}), ({O:3d}, {kD:2d}, {kH:2d}, {kW:2d}, {C:3d}), "
+                f"{strides}, {padding}, {groups:6d}, "
+                f"{100. * diff:+6.1f}%  "
+                f"{mlx_peak:8.1f}  {mlx_act:7.1f}  {pt_cur:7.1f}  {pt_drv:7.1f}"
+            )
@@ -1,5 +1,6 @@
 import math
 import os
+import platform
 import subprocess
 import time
 from copy import copy
@@ -17,9 +18,6 @@ RESULTS_DIR = "./results"
 if not os.path.isdir(RESULTS_DIR):
    os.mkdir(RESULTS_DIR)

-DEVICE_NAME = subprocess.check_output(["sysctl", "-n", "machdep.cpu.brand_string"])
-DEVICE_NAME = DEVICE_NAME.decode("utf-8").strip("\n")
-
 TORCH_DEVICE = torch.device(
    "mps"
    if torch.backends.mps.is_available()
@@ -27,11 +25,36 @@ TORCH_DEVICE = torch.device(
 )


+def get_device_name():
+    if TORCH_DEVICE.type == "cuda":
+        try:
+            out = subprocess.check_output(
+                ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
+                stderr=subprocess.DEVNULL,
+            )
+            return out.decode("utf-8").splitlines()[0].strip()
+        except Exception:
+            return "CUDA_GPU"
+    if TORCH_DEVICE.type == "mps":
+        try:
+            out = subprocess.check_output(
+                ["sysctl", "-n", "machdep.cpu.brand_string"],
+                stderr=subprocess.DEVNULL,
+            )
+            return out.decode("utf-8").strip()
+        except Exception:
+            return "Apple_Silicon"
+    return platform.processor() or platform.machine() or "CPU"
+
+
+DEVICE_NAME = get_device_name()
+
+
 N_WARMUP = 5
 N_ITER_BENCH = 50
 N_ITER_FUNC = 20

-VECTOR_LENGTHS = [4096 * (2**i) for i in range(10)]
+VECTOR_LENGTHS = [4096 * (2**i) for i in range(12)]
 MASK_DENSITIES = [0.01, 0.1, 0.25, 0.5]
 D_TYPES = ("float32", "float16")

@@ -202,9 +225,10 @@ def main():
        )
        output_path = os.path.join(
            RESULTS_DIR,
-            f"{DEVICE_NAME.replace(' ', '_')}_masked_scatter_{dtype}.pdf",
+            f"{DEVICE_NAME.replace(' ', '_')}_masked_scatter_{dtype}.png",
        )
        fig.savefig(output_path)
+        print(f"Saved benchmark image: {output_path}")
        plt.close(fig)


@@ -176,6 +176,8 @@ if __name__ == "__main__":
          (  1,  1024,  1024,       64,   32,     8),
          (  1,  2048,  2048,       64,   32,     8),
          (  1,  4096,  4096,       64,   32,     8),
+          (  1,  4096,  5000,       64,   32,     8),
+          (  1,  2048,  32121,      64,   32,     8),
    )

    shapes_80 = (
@@ -183,6 +185,8 @@ if __name__ == "__main__":
          (  1,  1024,  1024,       80,   32,     8),
          (  1,  2048,  2048,       80,   32,     8),
          (  1,  4096,  4096,       80,   32,     8),
+          (  1,  4096,  5000,       80,   32,     8),
+          (  1,  2048,  32121,      80,   32,     8),
    )

    shapes_128 = (
@@ -190,6 +194,8 @@ if __name__ == "__main__":
          (  1,  1024,  1024,      128,   32,     8),
          (  1,  2048,  2048,      128,   32,     8),
          (  1,  4096,  4096,      128,   32,     8),
+          (  1,  4096,  5000,      128,   32,     8),
+          (  1,  2048,  32121,     128,   32,     8),
    )
    # fmt: on

@@ -0,0 +1,209 @@
+# Copyright © 2026 Apple Inc.
+
+import argparse
+import time
+
+import mlx.core as mx
+import numpy as np
+
+MLX_DTYPES = {
+    "float16": mx.float16,
+    "bfloat16": mx.bfloat16,
+    "float32": mx.float32,
+}
+
+
+def parse_cases(cases):
+    parsed = []
+    for spec in cases.split(","):
+        m, n, k, s = [int(x) for x in spec.split("x")]
+        parsed.append((m, n, k, s))
+    return parsed
+
+
+def make_segments(k, num_segments, pattern, seed):
+    if pattern == "equal":
+        cuts = np.linspace(0, k, num_segments + 1, dtype=np.int64)
+    else:
+        rng = np.random.default_rng(seed)
+        cuts = rng.integers(0, k + 1, size=(num_segments - 1,), dtype=np.int64)
+        cuts = np.sort(cuts)
+        cuts = np.concatenate(([0], cuts, [k]))
+    return np.stack([cuts[:-1], cuts[1:]], axis=1).astype(np.uint32)
+
+
+def numpy_segmented_mm_ref(a, b, segments):
+    """Ground-truth reference in float64."""
+    out = []
+    for start, end in segments:
+        out.append(a[:, start:end] @ b[start:end, :])
+    return np.stack(out, axis=0)
+
+
+def mlx_segmented_mm_loop(a, b, segments):
+    """MLX loop-of-matmuls baseline."""
+    segments_list = segments.tolist()
+    out = []
+    for start, end in segments_list:
+        out.append(a[:, start:end] @ b[start:end, :])
+    return mx.stack(out, axis=0)
+
+
+def bench_mlx(a, b, segments, warmup, iters):
+    for _ in range(warmup):
+        y = mx.segmented_mm(a, b, segments)
+        mx.eval(y)
+    mx.synchronize()
+
+    start = time.perf_counter()
+    for _ in range(iters):
+        y = mx.segmented_mm(a, b, segments)
+        mx.eval(y)
+    mx.synchronize()
+    end = time.perf_counter()
+    return (end - start) * 1e3 / iters
+
+
+def bench_mlx_loop(a, b, segments, warmup, iters):
+    for _ in range(warmup):
+        y = mlx_segmented_mm_loop(a, b, segments)
+        mx.eval(y)
+    mx.synchronize()
+
+    start = time.perf_counter()
+    for _ in range(iters):
+        y = mlx_segmented_mm_loop(a, b, segments)
+        mx.eval(y)
+    mx.synchronize()
+    end = time.perf_counter()
+    return (end - start) * 1e3 / iters
+
+
+def print_table(headers, rows):
+    widths = [len(h) for h in headers]
+    for row in rows:
+        for i, cell in enumerate(row):
+            widths[i] = max(widths[i], len(cell))
+
+    def fmt_row(row):
+        return (
+            "| "
+            + " | ".join(f"{cell:<{widths[i]}}" for i, cell in enumerate(row))
+            + " |"
+        )
+
+    sep = "|-" + "-|-".join("-" * w for w in widths) + "-|"
+    print(fmt_row(headers))
+    print(sep)
+    for row in rows:
+        print(fmt_row(row))
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--cases",
+        default=(
+            "128x128x1024x16,"
+            "128x128x1024x32,"
+            "256x256x2048x16,"
+            "512x512x4096x32,"
+            "1024x1024x4096x32,"
+            "1024x1024x8192x64"
+        ),
+        help="Comma-separated MxNxKxS list.",
+    )
+    parser.add_argument(
+        "--dtype",
+        default="float32",
+        choices=["float16", "bfloat16", "float32"],
+    )
+    parser.add_argument("--warmup", type=int, default=10)
+    parser.add_argument("--iters", type=int, default=50)
+    parser.add_argument(
+        "--segments",
+        choices=["equal", "random"],
+        default="random",
+        help="Segment generation pattern.",
+    )
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--no-check", action="store_true")
+    args = parser.parse_args()
+
+    mlx_dtype = MLX_DTYPES[args.dtype]
+
+    print(
+        f"dtype={args.dtype} warmup={args.warmup} iters={args.iters} segments={args.segments}"
+    )
+
+    headers = [
+        "Case",
+        "MLX ms",
+        "Loop ms",
+        "Speedup",
+        "MLX err",
+        "Loop err",
+    ]
+    rows = []
+
+    cases = parse_cases(args.cases)
+    for idx, (m, n, k, s) in enumerate(cases):
+        rng = np.random.default_rng(args.seed + idx)
+        a_np = rng.standard_normal((m, k)).astype(np.float32)
+        b_np = rng.standard_normal((k, n)).astype(np.float32)
+        seg_np = make_segments(k, s, args.segments, args.seed + idx)
+
+        a_mx = mx.array(a_np, dtype=mlx_dtype)
+        b_mx = mx.array(b_np, dtype=mlx_dtype)
+        seg_mx = mx.array(seg_np, dtype=mx.uint32)
+        mx.eval(a_mx, b_mx, seg_mx)
+
+        mlx_err_str = ""
+        loop_err_str = ""
+        if not args.no_check:
+            y_mlx = mx.segmented_mm(a_mx, b_mx, seg_mx)
+            y_loop = mlx_segmented_mm_loop(a_mx, b_mx, seg_mx)
+            mx.eval(y_mlx, y_loop)
+
+            if args.dtype == "float32":
+                ref = numpy_segmented_mm_ref(
+                    a_np.astype(np.float64),
+                    b_np.astype(np.float64),
+                    seg_np.tolist(),
+                )
+                mlx_err = np.max(np.abs(np.array(y_mlx, dtype=np.float64) - ref))
+                loop_err = np.max(np.abs(np.array(y_loop, dtype=np.float64) - ref))
+            else:
+                a_mx_f32 = mx.array(a_np, dtype=mx.float32)
+                b_mx_f32 = mx.array(b_np, dtype=mx.float32)
+                ref = mx.segmented_mm(a_mx_f32, b_mx_f32, seg_mx)
+                mx.eval(ref)
+                mlx_err = float(mx.max(mx.abs(ref - y_mlx.astype(mx.float32))).item())
+                loop_err = float(mx.max(mx.abs(ref - y_loop.astype(mx.float32))).item())
+            mlx_err_str = f"{mlx_err:.2e}"
+            loop_err_str = f"{loop_err:.2e}"
+
+        t_mlx = bench_mlx(a_mx, b_mx, seg_mx, args.warmup, args.iters)
+        t_loop = bench_mlx_loop(a_mx, b_mx, seg_mx, args.warmup, args.iters)
+        ratio = t_loop / t_mlx if t_mlx > 0 else float("inf")
+        rows.append(
+            [
+                f"{m}x{n}x{k}x{s}",
+                f"{t_mlx:.3f}",
+                f"{t_loop:.3f}",
+                f"{ratio:.2f}x",
+                mlx_err_str,
+                loop_err_str,
+            ]
+        )
+
+    print_table(headers, rows)
+    if not args.no_check:
+        if args.dtype == "float32":
+            print("err: max|result - numpy_fp64_ref|")
+        else:
+            print("err: max|result - own_fp32_result|")
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,109 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import argparse
+
+import mlx.core as mx
+import torch
+from time_utils import measure_runtime
+
+
+def benchmark_slice_update_mlx(dst_shape, slice_shape, slice_range, dtype, iters=10):
+    def slice_update(arguments):
+        for i in range(iters):
+            arguments["dst"] = (
+                arguments["dst"].at[slice_range].add(arguments["updates"])
+            )
+        mx.eval(arguments)
+
+    dtype = getattr(mx, dtype)
+    arguments = {
+        "dst": mx.random.normal(dst_shape).astype(dtype),
+        "updates": mx.random.normal(slice_shape).astype(dtype),
+    }
+
+    runtime = measure_runtime(slice_update, arguments=arguments)
+    bytes_processed = (
+        arguments["dst"][slice_range].nbytes * 2 + arguments["updates"].nbytes
+    ) * iters
+    bandwidth_gb_s = bytes_processed / runtime / 1e6
+    return runtime, bandwidth_gb_s
+
+
+def benchmark_slice_update_torch(
+    dst_shape, slice_shape, slice_range, device, dtype, iters=10
+):
+    def slice_update(dst, updates, slice_range):
+        for i in range(iters):
+            dst[slice_range] = dst[slice_range] + updates
+        if device == torch.device("mps"):
+            torch.mps.synchronize()
+
+    dtype = getattr(torch, dtype)
+    updates = torch.randn(slice_shape, dtype=dtype).to(device)
+    dst = torch.randn(dst_shape, dtype=dtype).to(device)
+
+    runtime = measure_runtime(
+        slice_update, dst=dst, updates=updates, slice_range=slice_range
+    )
+    bytes_processed = (dst[slice_range].nbytes * 2 + updates.nbytes) * iters
+    bandwidth_gb_s = bytes_processed / runtime / 1e6
+    return runtime, bandwidth_gb_s
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser("Slice update benchmarks.")
+    parser.add_argument("--cpu", action="store_true", help="Use the CPU.")
+    args = parser.parse_args()
+
+    if args.cpu:
+        mx.set_default_device(mx.cpu)
+        device = torch.device("cpu")
+    elif torch.mps.is_available():
+        device = torch.device("mps")
+    elif torch.cuda.is_available():
+        device = torch.device("cuda")
+    else:
+        raise ValueError()
+
+    dtypes = ["float32", "bfloat16"]
+
+    test_cases = [
+        ((10_000_000,), slice(0, 1_000_000), (1_000_000,)),
+        ((100_000,), slice(10_000, 20_000), (10_000,)),
+        ((1000, 64), slice(100, 200), (100, 64)),
+        ((100, 100, 64), slice(20, 40), (20, 100, 64)),
+        (
+            (2048, 2048, 128),
+            (slice(500, 1500), slice(200, 1200), slice(32, 96)),
+            (1000, 1000, 64),
+        ),
+        (
+            (2048, 2048, 128),
+            (slice(1800, 1850), slice(100, 200), slice(64, 128)),
+            (50, 100, 64),
+        ),
+        (
+            (2048, 2048, 128),
+            (slice(1000, 1010), slice(1000, 1010), slice(64, 128)),
+            (10, 10, 64),
+        ),
+    ]
+
+    print(
+        f"{'Dtype':<12} {'Dst Shape':<25} {'Update Shape':<20} "
+        f"{'MLX (ms)':<12} {'MLX GB/s':<12} {'Torch (ms)':<12} {'Torch GB/s':<12}"
+    )
+    print("-" * 110)
+
+    for dtype in dtypes:
+        for dst_shape, slice_range, update_shape in test_cases:
+            mlx_time, mlx_bw = benchmark_slice_update_mlx(
+                dst_shape, update_shape, slice_range, dtype
+            )
+            torch_time, torch_bw = benchmark_slice_update_torch(
+                dst_shape, update_shape, slice_range, device, dtype
+            )
+            print(
+                f"{dtype:<12} {str(dst_shape):<25} {str(update_shape):<20} "
+                f"{mlx_time:<12.3f} {mlx_bw:<12.2f} {torch_time:<12.3f} {torch_bw:<12.2f}"
+            )
@@ -38,3 +38,17 @@ the docs. Then force add the `build/html` directory:
 `git add -f build/html`

 Commit and push the changes to the `gh-pages` branch.
+
+## Doc Development Setup
+
+To enable live refresh of docs while writing:
+
+Install sphinx autobuild
+```
+pip install sphinx-autobuild
+```
+
+Run auto build on docs/src folder
+```
+sphinx-autobuild ./src ./build/html
+```
@@ -0,0 +1,36 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" width="433.19" height="139.72" viewBox="0 0 433.19 139.72">
+<defs>
+<g>
+<g id="glyph-0-0">
+<path d="M 9.6875 -171.1875 L 188.609375 -171.1875 L 188.609375 -188.8125 L 9.6875 -188.8125 Z M 9.6875 -127.421875 L 188.609375 -127.421875 L 188.609375 -145.046875 L 9.6875 -145.046875 Z M 9.6875 -83.5625 L 188.609375 -83.5625 L 188.609375 -101.1875 L 9.6875 -101.1875 Z M 9.6875 -39.796875 L 188.609375 -39.796875 L 188.609375 -57.421875 L 9.6875 -57.421875 Z M 9.6875 4.0625 L 188.609375 4.0625 L 188.609375 -13.5625 L 9.6875 -13.5625 Z M 9.6875 47.828125 L 188.609375 47.828125 L 188.609375 30.203125 L 9.6875 30.203125 Z M 9.6875 47.828125 "/>
+</g>
+<g id="glyph-0-1">
+<path d="M 13.9375 0 L 42.3125 0 L 42.3125 -91.796875 L 43.671875 -91.796875 L 78.71875 0 L 97.984375 0 L 133.03125 -91.796875 L 134.390625 -91.796875 L 134.390625 0 L 162.765625 0 L 162.765625 -139.71875 L 125.96875 -139.71875 L 88.984375 -42.015625 L 87.8125 -42.015625 L 50.734375 -139.71875 L 13.9375 -139.71875 Z M 13.9375 0 "/>
+</g>
+<g id="glyph-0-2">
+<path d="M 13.9375 0 L 106.21875 0 L 106.21875 -26.046875 L 45.984375 -26.046875 L 45.984375 -139.71875 L 13.9375 -139.71875 Z M 13.9375 0 "/>
+</g>
+<g id="glyph-0-3">
+<path d="M 6.296875 0 L 40.5625 0 L 69.515625 -47.34375 L 70.484375 -47.34375 L 99.625 0 L 135.84375 0 L 91.109375 -69.8125 L 91.109375 -70.296875 L 136.328125 -139.71875 L 100.703125 -139.71875 L 73 -90.71875 L 71.84375 -90.71875 L 43.953125 -139.71875 L 6.484375 -139.71875 L 49.96875 -70.6875 L 49.96875 -70.296875 Z M 6.296875 0 "/>
+</g>
+</g>
+<clipPath id="clip-0">
+<path clip-rule="nonzero" d="M 13 0 L 283 0 L 283 139.71875 L 13 139.71875 Z M 13 0 "/>
+</clipPath>
+<clipPath id="clip-1">
+<path clip-rule="nonzero" d="M 296 0 L 427 0 L 427 139.71875 L 296 139.71875 Z M 296 0 "/>
+</clipPath>
+</defs>
+<g clip-path="url(#clip-0)">
+<g fill="rgb(0%, 0%, 0%)" fill-opacity="1">
+<use xlink:href="#glyph-0-1" x="0" y="139.72"/>
+<use xlink:href="#glyph-0-2" x="176.682092" y="139.72"/>
+</g>
+</g>
+<g clip-path="url(#clip-1)">
+<g fill="rgb(82.998657%, 82.998657%, 82.998657%)" fill-opacity="1">
+<use xlink:href="#glyph-0-3" x="290.57" y="139.72"/>
+</g>
+</g>
+</svg>
@@ -0,0 +1,36 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" width="433.19" height="139.72" viewBox="0 0 433.19 139.72">
+<defs>
+<g>
+<g id="glyph-0-0">
+<path d="M 9.6875 -171.1875 L 188.609375 -171.1875 L 188.609375 -188.8125 L 9.6875 -188.8125 Z M 9.6875 -127.421875 L 188.609375 -127.421875 L 188.609375 -145.046875 L 9.6875 -145.046875 Z M 9.6875 -83.5625 L 188.609375 -83.5625 L 188.609375 -101.1875 L 9.6875 -101.1875 Z M 9.6875 -39.796875 L 188.609375 -39.796875 L 188.609375 -57.421875 L 9.6875 -57.421875 Z M 9.6875 4.0625 L 188.609375 4.0625 L 188.609375 -13.5625 L 9.6875 -13.5625 Z M 9.6875 47.828125 L 188.609375 47.828125 L 188.609375 30.203125 L 9.6875 30.203125 Z M 9.6875 47.828125 "/>
+</g>
+<g id="glyph-0-1">
+<path d="M 13.9375 0 L 42.3125 0 L 42.3125 -91.796875 L 43.671875 -91.796875 L 78.71875 0 L 97.984375 0 L 133.03125 -91.796875 L 134.390625 -91.796875 L 134.390625 0 L 162.765625 0 L 162.765625 -139.71875 L 125.96875 -139.71875 L 88.984375 -42.015625 L 87.8125 -42.015625 L 50.734375 -139.71875 L 13.9375 -139.71875 Z M 13.9375 0 "/>
+</g>
+<g id="glyph-0-2">
+<path d="M 13.9375 0 L 106.21875 0 L 106.21875 -26.046875 L 45.984375 -26.046875 L 45.984375 -139.71875 L 13.9375 -139.71875 Z M 13.9375 0 "/>
+</g>
+<g id="glyph-0-3">
+<path d="M 6.296875 0 L 40.5625 0 L 69.515625 -47.34375 L 70.484375 -47.34375 L 99.625 0 L 135.84375 0 L 91.109375 -69.8125 L 91.109375 -70.296875 L 136.328125 -139.71875 L 100.703125 -139.71875 L 73 -90.71875 L 71.84375 -90.71875 L 43.953125 -139.71875 L 6.484375 -139.71875 L 49.96875 -70.6875 L 49.96875 -70.296875 Z M 6.296875 0 "/>
+</g>
+</g>
+<clipPath id="clip-0">
+<path clip-rule="nonzero" d="M 13 0 L 283 0 L 283 139.71875 L 13 139.71875 Z M 13 0 "/>
+</clipPath>
+<clipPath id="clip-1">
+<path clip-rule="nonzero" d="M 296 0 L 427 0 L 427 139.71875 L 296 139.71875 Z M 296 0 "/>
+</clipPath>
+</defs>
+<g clip-path="url(#clip-0)">
+<g fill="rgb(100%, 100%, 100%)" fill-opacity="1">
+<use xlink:href="#glyph-0-1" x="0" y="139.72"/>
+<use xlink:href="#glyph-0-2" x="176.682092" y="139.72"/>
+</g>
+</g>
+<g clip-path="url(#clip-1)">
+<g fill="rgb(56.999207%, 56.999207%, 56.999207%)" fill-opacity="1">
+<use xlink:href="#glyph-0-3" x="290.57" y="139.72"/>
+</g>
+</g>
+</svg>
@@ -404,7 +404,7 @@ below.
        auto kernel = d.get_kernel(kname, lib);

        // Prepare to encode kernel
-        auto& compute_encoder = d.get_command_encoder(s.index);
+        auto& compute_encoder = mx::metal::get_command_encoder(s);
        compute_encoder.set_compute_pipeline_state(kernel);

        // Kernel parameters are registered with buffer indices corresponding to
@@ -448,7 +448,7 @@ We can now call the :meth:`axpby` operation on both the CPU and the GPU!

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

 .. toctree::
-   :caption: Usage 
+   :caption: Usage
   :maxdepth: 1

   usage/quick_start
@@ -54,6 +54,8 @@ are the CPU and GPU.
   examples/linear_regression
   examples/mlp
   examples/llama-inference
+   examples/data_parallelism
+   examples/tensor_parallelism

 .. toctree::
   :caption: Python API Reference
@@ -76,6 +78,7 @@ are the CPU and GPU.
   python/optimizers
   python/distributed
   python/tree_utils
+   python/printoptions

 .. toctree::
   :caption: C++ API Reference
@@ -15,7 +15,7 @@ silicon computer is

 To install from PyPI your system must meet the following requirements:

- Using an M series chip (Apple silicon)
+- Using `Apple silicon <https://support.apple.com/en-us/116943>`_
 - Using a native Python >= 3.10
 - macOS >= 14.0

@@ -83,6 +83,7 @@ Build from source
 Build Requirements
 ^^^^^^^^^^^^^^^^^^

+- ``libblas-dev``, ``liblapack-dev``, and ``liblapacke-dev`` (Linux)
 - A C++ compiler with C++20 support (e.g. Clang >= 15.0)
 - `cmake <https://cmake.org/>`_ -- version 3.25 or later, and ``make``
 - Xcode >= 15.0 and macOS SDK >= 14.0
@@ -14,8 +14,10 @@ Devices and Streams
   set_default_device
   default_stream
   new_stream
+   new_thread_local_stream
   set_default_stream
   stream
   synchronize
+   clear_streams
   device_count
   device_info
@@ -20,5 +20,7 @@ FFT
  irfft2
  rfftn
  irfftn
+  fftfreq
+  rfftfreq
  fftshift
  ifftshift
@@ -14,6 +14,7 @@ Linear Algebra
    cholesky
    cholesky_inv
    cross
+    det
    qr
    svd
    eigvals
@@ -23,5 +24,6 @@ Linear Algebra
    lu
    lu_factor
    pinv
+    slogdet
    solve
    solve_triangular
@@ -175,6 +175,7 @@ In detail:
   value_and_grad
   quantize
   average_gradients
+   fsdp_apply_gradients

 .. toctree::

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

   ALiBi
+   AllToShardedLinear
   AvgPool1d
   AvgPool2d
   AvgPool3d
@@ -46,8 +47,10 @@ Layers
   Mish
   MultiHeadAttention
   PReLU
+   QuantizedAllToShardedLinear
   QuantizedEmbedding
   QuantizedLinear
+   QuantizedShardedToAllLinear
   RMSNorm
   ReLU
   ReLU2
@@ -56,6 +59,7 @@ Layers
   RoPE
   SELU
   Sequential
+   ShardedToAllLinear
   Sigmoid
   SiLU
   SinusoidalPositionalEncoding
@@ -0,0 +1,12 @@
+Print Options
+===============
+
+.. currentmodule:: mlx.core
+
+.. autosummary::
+  :toctree: _autosummary
+
+  PrintOptions
+  set_printoptions
+  printoptions
+  get_printoptions
@@ -117,89 +117,11 @@ The following examples aim to clarify the backend initialization logic in MLX:
    world_ring = mx.distributed.init(backend="ring")
    world_any = mx.distributed.init()  # same as MPI because it was initialized first!

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

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

 .. _ring_section:

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


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

@@ -90,10 +90,7 @@ PyTorch supports the buffer protocol, but it requires an explicit

  a = mx.arange(3)
  b = torch.tensor(memoryview(a))
-  c = mx.array(b.numpy())
-
-Conversion from PyTorch tensors back to arrays must be done via intermediate
-NumPy arrays with ``numpy()``.
+  c = mx.array(b)

 JAX
 ---
@@ -192,7 +192,7 @@ void Axpby::eval_gpu(
  auto kernel = d.get_kernel(kname, lib);

  // Prepare to encode kernel
-  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto& compute_encoder = mx::metal::get_command_encoder(s);
  compute_encoder.set_compute_pipeline_state(kernel);

  // Kernel parameters are registered with buffer indices corresponding to
@@ -3,6 +3,6 @@ requires = [
  "setuptools>=42",
  "cmake>=3.25",
  "mlx>=0.18.0",
-  "nanobind==2.10.2",
+  "nanobind==2.12.0",
 ]
 build-backend = "setuptools.build_meta"
@@ -1,4 +1,4 @@
 setuptools>=42
 cmake>=3.25
-mlx>=0.21.0
-nanobind==2.10.2
+mlx>=0.31.2
+nanobind==2.12.0
@@ -14,6 +14,7 @@ target_sources(
          ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/random.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/scheduler.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/stream.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/transforms.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/linalg.cpp
@@ -32,10 +33,11 @@ set_target_properties(
             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)
+# Define MLX_EXPORT for shared libraries, MLX_STATIC for static libraries.
+set_target_properties(mlx PROPERTIES DEFINE_SYMBOL MLX_EXPORT)
+if(BUILD_SHARED_LIBS)
+  target_compile_definitions(mlx_version PUBLIC MLX_EXPORT)
+else()
  target_compile_definitions(mlx PUBLIC MLX_STATIC)
  target_compile_definitions(mlx_version PUBLIC MLX_STATIC)
 endif()
@@ -49,20 +51,20 @@ endif()

 if(MSVC)
  # Some of CUDA's headers include windows.h, which defines min/max macros.
-  target_compile_definitions(mlx PRIVATE NOMINMAX)
+  target_compile_definitions(mlx PRIVATE NOMINMAX WIN32_LEAN_AND_MEAN)
+  # Unicode support in fmt does not compile in .cu files.
+  target_compile_definitions(mlx PRIVATE FMT_UNICODE=0)
  # Disable some MSVC warnings to speed up compilation.
  target_compile_options(
    mlx
-    PUBLIC $<$<COMPILE_LANGUAGE:CXX>:/wd4068
-           /wd4244
-           /wd4267
-           /wd4700
-           /wd4804>
-           $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/wd4068
-           -Xcompiler=/wd4244
-           -Xcompiler=/wd4267
-           -Xcompiler=/wd4700
-           -Xcompiler=/wd4804>)
+    PUBLIC $<$<COMPILE_LANGUAGE:CXX>:/wd4244 /wd4267>
+    PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/wd4068
+            /wd4146
+            /wd4700
+            /wd4804
+            /wd4805>
+            $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/wd4244
+            -Xcompiler=/wd4267>)
  # Enable /bigobj for heavily templated code (e.g., binary.cpp) that exceeds
  # the default 65,535 section limit in COFF object files.
  target_compile_options(
@@ -134,6 +134,7 @@ bool array::is_available() const {
  } else if (
      status() == Status::evaluated &&
      (!event().valid() || event().is_signaled())) {
+    detach_event();
    set_status(Status::available);
    return true;
  }
@@ -489,10 +489,10 @@ class MLX_API array {
    int64_t offset{0};

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

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

    std::vector<array> inputs;
    // An array to keep track of the siblings from a multi-output
@@ -0,0 +1,14 @@
+// Copyright © 2026 Apple Inc.
+
+namespace mlx::core {
+
+inline constexpr short get_pack_factor(int bits, int wsize = 8) {
+  return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits);
+}
+
+inline constexpr short get_bytes_per_pack(int bits, int wsize = 8) {
+  bool power_of_2_bits = (bits & (bits - 1)) == 0;
+  return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3);
+}
+
+} // namespace mlx::core
@@ -116,6 +116,39 @@ struct ContiguousIterator {
    loc += strides_[i];
  }

+  void step(int64_t s) {
+    int dims = shape_.size();
+    if (dims == 0) {
+      return;
+    }
+    int i = dims - 1;
+    while (s > 0) {
+      if (shape_[i] - pos_[i] > 1) {
+        int steps = static_cast<int>(
+            std::min(static_cast<int64_t>(shape_[i] - pos_[i] - 1), s));
+        pos_[i] += steps;
+        loc += strides_[i] * steps;
+        s -= steps;
+      } else {
+        while (pos_[i] == (shape_[i] - 1) && i > 0) {
+          pos_[i] = 0;
+          loc -= (shape_[i] - 1) * strides_[i];
+          i--;
+        }
+        pos_[i]++;
+        loc += strides_[i];
+        s--;
+      }
+    }
+  }
+
+  int64_t contiguous_suffix() {
+    if (shape_.size() == 0) {
+      return 0;
+    }
+    return (strides_.back() == 1) ? shape_.back() : 0;
+  }
+
  void seek(int64_t n) {
    loc = 0;
    for (int i = shape_.size() - 1; i >= 0; --i) {
@@ -6,8 +6,6 @@
 #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>
@@ -4,11 +4,14 @@
 #include <cmath>

 #include "mlx/allocator.h"
-#include "mlx/primitives.h"
-
 #include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/binary.h"
+#include "mlx/backend/cpu/binary_ops.h"
 #include "mlx/backend/cpu/copy.h"
 #include "mlx/backend/cpu/encoder.h"
+#include "mlx/backend/cpu/slicing.h"
+#include "mlx/dtype_utils.h"
+#include "mlx/primitives.h"

 namespace mlx::core {

@@ -788,7 +791,7 @@ void MaskedScatter::eval_cpu(const std::vector<array>& inputs, array& out) {
  auto& mask = inputs[1];
  auto& src = inputs[2];

-  // Copy src into out (copy allocates memory for out)
+  // Copy dst into out (copy allocates memory for out)
  auto ctype =
      dst.flags().row_contiguous ? CopyType::Vector : CopyType::General;
  copy_cpu(dst, out, ctype, stream());
@@ -851,4 +854,128 @@ void MaskedScatter::eval_cpu(const std::vector<array>& inputs, array& out) {
  });
 }

+template <typename T, typename Op>
+void slice_update_impl(
+    array& out,
+    const array& upd,
+    int64_t data_offset,
+    const Strides& out_strides) {
+  ContiguousIterator out_it(upd.shape(), out_strides, upd.ndim());
+  ContiguousIterator upd_it(upd);
+  Op op;
+
+  constexpr int SIMD_START = 32;
+
+  T* out_ptr = out.data<T>() + data_offset;
+  const T* upd_ptr = upd.data<T>();
+  int64_t size = upd.size();
+  int64_t suffix = out_it.contiguous_suffix();
+
+  if (upd.data_size() == 1) {
+    if (suffix >= SIMD_START) {
+      for (int64_t i = 0; i < size; i += suffix) {
+        VectorScalar<Op>{}(
+            out_ptr + out_it.loc, upd_ptr, out_ptr + out_it.loc, suffix);
+        out_it.step(suffix);
+      }
+    } else {
+      T update = upd_ptr[0];
+      for (int64_t i = 0; i < size; i++) {
+        out_ptr[out_it.loc] = op(out_ptr[out_it.loc], update);
+        out_it.step();
+      }
+    }
+  } else if (suffix == upd_it.contiguous_suffix() && suffix >= SIMD_START) {
+    for (int64_t i = 0; i < size; i += suffix) {
+      VectorVector<Op>{}(
+          out_ptr + out_it.loc,
+          upd_ptr + upd_it.loc,
+          out_ptr + out_it.loc,
+          suffix);
+      out_it.step(suffix);
+      upd_it.step(suffix);
+    }
+  } else {
+    for (int64_t i = 0; i < size; i++) {
+      out_ptr[out_it.loc] = op(out_ptr[out_it.loc], upd_ptr[upd_it.loc]);
+      out_it.step();
+      upd_it.step();
+    }
+  }
+}
+
+void SliceUpdate::eval_cpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 2);
+  if (out.size() == 0) {
+    out.set_data(allocator::malloc(0));
+    return;
+  }
+
+  auto& in = inputs[0];
+  auto& upd = inputs[1];
+
+  if (upd.size() == 0) {
+    out.copy_shared_buffer(in);
+    return;
+  }
+
+  // Check if materialization is needed
+  auto ctype = in.flags().contiguous && in.size() == in.data_size()
+      ? CopyType::Vector
+      : CopyType::General;
+  copy_cpu(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype, stream());
+
+  // Calculate out strides, initial offset and if copy needs to be made
+  auto [data_offset, out_strides] =
+      prepare_slice(out, start_indices_, strides_);
+
+  // Do copy
+  if (reduce_type_ == SliceUpdate::None) {
+    copy_cpu_inplace(
+        /* const array& src = */ upd,
+        /* array& dst = */ out,
+        /* const std::vector<int>& data_shape = */ upd.shape(),
+        /* const std::vector<stride_t>& i_strides = */ upd.strides(),
+        /* const std::vector<stride_t>& o_strides = */ out_strides,
+        /* int64_t i_offset = */ 0,
+        /* int64_t o_offset = */ data_offset,
+        /* CopyType ctype = */ CopyType::GeneralGeneral,
+        stream());
+    return;
+  }
+
+  auto& encoder = cpu::get_command_encoder(stream());
+  encoder.set_input_array(upd);
+  encoder.set_output_array(out);
+  encoder.dispatch([upd = array::unsafe_weak_copy(upd),
+                    out = array::unsafe_weak_copy(out),
+                    data_offset = data_offset,
+                    out_strides = std::move(out_strides),
+                    reduce_type = reduce_type_]() mutable {
+    dispatch_all_types(out.dtype(), [&](auto type_tag) {
+      using T = MLX_GET_TYPE(type_tag);
+      switch (reduce_type) {
+        case SliceUpdate::Sum:
+          slice_update_impl<T, detail::Add>(out, upd, data_offset, out_strides);
+          break;
+        case SliceUpdate::Prod:
+          slice_update_impl<T, detail::Multiply>(
+              out, upd, data_offset, out_strides);
+          break;
+        case SliceUpdate::Max:
+          slice_update_impl<T, detail::Maximum>(
+              out, upd, data_offset, out_strides);
+          break;
+        case SliceUpdate::Min:
+          slice_update_impl<T, detail::Minimum>(
+              out, upd, data_offset, out_strides);
+          break;
+        case SliceUpdate::None:
+          // Should never be here
+          break;
+      }
+    });
+  });
+}
+
 } // namespace mlx::core
@@ -67,11 +67,10 @@ void luf_impl(
              /* ipiv */ reinterpret_cast<int*>(pivots_ptr),
              /* info */ &info);

-          if (info != 0) {
+          if (info < 0) {
            std::stringstream ss;
            ss << "[LUF::eval_cpu] sgetrf_ failed with code " << info
-               << ((info > 0) ? " because matrix is singular"
-                              : " because argument had an illegal value");
+               << " because argument had an illegal value";
            throw std::runtime_error(ss.str());
          }

@@ -398,44 +398,6 @@ void DynamicSliceUpdate::eval_cpu(
  }
 }

-void SliceUpdate::eval_cpu(const std::vector<array>& inputs, array& out) {
-  assert(inputs.size() == 2);
-  if (out.size() == 0) {
-    out.set_data(allocator::malloc(0));
-    return;
-  }
-
-  auto& in = inputs[0];
-  auto& upd = inputs[1];
-
-  if (upd.size() == 0) {
-    out.copy_shared_buffer(in);
-    return;
-  }
-
-  // Check if materialization is needed
-  auto ctype = in.flags().contiguous && in.size() == in.data_size()
-      ? CopyType::Vector
-      : CopyType::General;
-  copy_cpu(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype, stream());
-
-  // Calculate out strides, initial offset and if copy needs to be made
-  auto [data_offset, out_strides] =
-      prepare_slice(out, start_indices_, strides_);
-
-  // Do copy
-  copy_cpu_inplace(
-      /* const array& src = */ upd,
-      /* array& dst = */ out,
-      /* const std::vector<int>& data_shape = */ upd.shape(),
-      /* const std::vector<stride_t>& i_strides = */ upd.strides(),
-      /* const std::vector<stride_t>& o_strides = */ out_strides,
-      /* int64_t i_offset = */ 0,
-      /* int64_t o_offset = */ data_offset,
-      /* CopyType ctype = */ CopyType::GeneralGeneral,
-      stream());
-}
-
 void View::eval_cpu(const std::vector<array>& inputs, array& out) {
  assert(inputs.size() == 1);
  auto& in = inputs[0];
@@ -1,5 +1,6 @@
 // Copyright © 2023 Apple Inc.

+#include "mlx/backend/common/quantized.h"
 #include "mlx/backend/common/unary.h"
 #include "mlx/backend/cpu/copy.h"
 #include "mlx/backend/cpu/encoder.h"
@@ -60,15 +61,6 @@ static inline T dequantize_scale(uint8_t s) {
  }
 }

-inline constexpr short get_pack_factor(int bits, int wsize = 8) {
-  return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits);
-}
-
-inline constexpr short get_bytes_per_pack(int bits, int wsize = 8) {
-  auto power_of_2_bits = (bits & (bits - 1)) == 0;
-  return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3);
-}
-
 template <typename T, int bits>
 void extract_bits(const uint8_t* w_in, T* w_out) {
  static_assert(bits == 3 || bits == 5 || bits == 6);
@@ -29,6 +29,14 @@ struct Simd<T, 1> {
  Simd(Simd<U, 1> v) : value(v.value) {}
  template <typename U>
  Simd(U v) : value(v) {}
+
+  T operator[](int) const {
+    return value;
+  }
+
+  T& operator[](int) {
+    return value;
+  }
 };

 template <typename T, int N>
@@ -15,10 +15,14 @@ namespace mlx::core {

 namespace {

+template <typename T>
+inline constexpr bool is_floating_v = std::is_floating_point_v<T> ||
+    std::is_same_v<T, float16_t> || std::is_same_v<T, bfloat16_t>;
+
 // NaN-aware comparator that places NaNs at the end
 template <typename T>
 bool nan_aware_less(T a, T b) {
-  if constexpr (std::is_floating_point_v<T> || std::is_same_v<T, complex64_t>) {
+  if constexpr (is_floating_v<T> || std::is_same_v<T, complex64_t>) {
    if (std::isnan(a))
      return false;
    if (std::isnan(b))
@@ -103,11 +107,11 @@ struct StridedIterator {
    return *this;
  }

-  StridedIterator operator+(difference_type diff) {
+  StridedIterator operator+(difference_type diff) const {
    return StridedIterator(ptr_, stride_, diff);
  }

-  StridedIterator operator-(difference_type diff) {
+  StridedIterator operator-(difference_type diff) const {
    return StridedIterator(ptr_, stride_, -diff);
  }

@@ -198,7 +202,7 @@ void argsort(const array& in, array& out, int axis) {
      auto v2 = data_ptr[b * in_stride];

      // Handle NaNs (place them at the end)
-      if (std::is_floating_point<T>::value) {
+      if constexpr (is_floating_v<T>) {
        if (std::isnan(v1))
          return false;
        if (std::isnan(v2))
@@ -299,7 +303,7 @@ void argpartition(const array& in, array& out, int axis, int kth) {
      auto v2 = data_ptr[b * in_stride];

      // Handle NaNs (place them at the end)
-      if (std::is_floating_point<T>::value) {
+      if constexpr (is_floating_v<T>) {
        if (std::isnan(v1))
          return false;
        if (std::isnan(v2))
@@ -155,11 +155,18 @@ struct FromFP8 {
  template <int N>
  Simd<float, N> operator()(Simd<uint8_t, N> x) {
    auto v = Simd<uint16_t, N>(x & 127) << 7;
-    auto converted = *(Simd<float16_t, N>*)(&v);
-    converted = converted * 256.0;
+    Simd<float, N> out;
+    if constexpr (simd::max_size<float16_t> >= N) {
+      auto converted = *(Simd<float16_t, N>*)(&v);
+      out = converted * 256.0;
+    } else {
+      for (int i = 0; i < N; ++i) {
+        auto converted = *(float16_t*)(&v[i]);
+        out[i] = converted * 256.0;
+      }
+    }
    auto sign = Simd<bool, N>(x & 128);
-    Simd<float, N> out = select(sign, -converted, converted);
-    return out;
+    return select(sign, -out, out);
  }
  float operator()(uint8_t x) {
    return (*this)(Simd<uint8_t, 1>(x)).value;
@@ -26,10 +26,14 @@ target_sources(
          ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cu
          ${CMAKE_CURRENT_SOURCE_DIR}/eval.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/event.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/fft.cu
          ${CMAKE_CURRENT_SOURCE_DIR}/fence.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/gemms/block_mask.cu
          ${CMAKE_CURRENT_SOURCE_DIR}/gemms/gemv.cu
          ${CMAKE_CURRENT_SOURCE_DIR}/gemms/cublas_gemm.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/gemms/gather_gemm.cu
          ${CMAKE_CURRENT_SOURCE_DIR}/gemms/grouped_gemm_unaligned.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cu
          ${CMAKE_CURRENT_SOURCE_DIR}/jit_module.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
          ${CMAKE_CURRENT_SOURCE_DIR}/kernel_utils.cu
@@ -56,7 +60,6 @@ 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
@@ -64,6 +67,7 @@ target_sources(
          ${CMAKE_CURRENT_SOURCE_DIR}/worker.cpp)

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

 # fp4 is not available on < 12.8
@@ -116,6 +120,16 @@ target_compile_options(mlx
 target_compile_options(
  mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>")

+# Ignore warnings from CUTLASS.
+target_compile_options(
+  mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe="--diag_suppress=2908,2361">)
+
+if(NOT MSVC)
+  # Required for generating optimized CUTLASS code.
+  target_compile_options(
+    mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-fno-strict-aliasing>")
+endif()
+
 # Suppress nvcc warnings on C++ headers.
 target_compile_options(
  mlx
@@ -140,12 +154,11 @@ if(NOT DEFINED MLX_CUDA_ARCHITECTURES)
    COMMAND __nvcc_device_query
    OUTPUT_VARIABLE MLX_CUDA_ARCHITECTURES
    OUTPUT_STRIP_TRAILING_WHITESPACE)
-  set(UPGRADABLE_ARCHITECTURES "90;100;121")
  if(MLX_CUDA_ARCHITECTURES STREQUAL "")
    message(
      FATAL_ERROR
        "Can not get native CUDA arch, must set MLX_CUDA_ARCHITECTURES")
-  elseif(MLX_CUDA_ARCHITECTURES IN_LIST UPGRADABLE_ARCHITECTURES)
+  elseif(MLX_CUDA_ARCHITECTURES GREATER_EQUAL 90)
    # Use arch-specific compute capability whenever possible.
    set(MLX_CUDA_ARCHITECTURES "${MLX_CUDA_ARCHITECTURES}a")
  endif()
@@ -154,6 +167,12 @@ message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}")
 set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES
                                     "${MLX_CUDA_ARCHITECTURES}")

+# Skip Hopper-only kernels when not building for sm90a.
+if(("90a" IN_LIST MLX_CUDA_ARCHITECTURES) OR ("90a-real" IN_LIST
+                                              MLX_CUDA_ARCHITECTURES))
+  target_compile_definitions(mlx PRIVATE MLX_CUDA_SM90A_ENABLED)
+endif()
+
 # Search CUDA libs from installed python packages.
 if(WIN32)
  # Resolve paths of unfound DLL at runtime.
@@ -234,6 +253,9 @@ target_include_directories(mlx PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
 # Use cublasLt.
 target_link_libraries(mlx PRIVATE CUDA::cublasLt)

+# Use cuFFT.
+target_link_libraries(mlx PRIVATE CUDA::cufft)
+
 # Use NVRTC and driver APIs.
 target_link_libraries(mlx PRIVATE CUDA::nvrtc CUDA::cuda_driver)

@@ -257,7 +279,7 @@ target_link_libraries(mlx PRIVATE CUDNN::cudnn_all)
 FetchContent_Declare(
  cutlass
  GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
-  GIT_TAG v4.3.5
+  GIT_TAG v4.4.2
  GIT_SHALLOW TRUE
  SOURCE_SUBDIR include EXCLUDE_FROM_ALL)
 FetchContent_MakeAvailable(cutlass)
@@ -12,6 +12,8 @@
 #include <fmt/format.h>

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

 namespace mlx::core {

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

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

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

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

@@ -65,7 +127,7 @@ SmallSizePool::SmallSizePool() {
 }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  CudaAllocator();
  friend CudaAllocator& allocator();
@@ -77,7 +82,7 @@ class CudaAllocator : public allocator::Allocator {
  BufferCache<CudaBuffer> buffer_cache_;
  size_t active_memory_{0};
  size_t peak_memory_{0};
-  std::vector<cudaStream_t> free_streams_;
+  std::vector<CudaStream> free_streams_;
  std::vector<cudaMemPool_t> mem_pools_;
  SmallSizePool scalar_pool_;
 };
@@ -56,7 +56,6 @@ void Arange::eval_gpu(const std::vector<array>& inputs, array& out) {
          cu::arange<OutType, IdxT, N_WRITES>,
          num_blocks,
          block_dims,
-          0,
          gpu_ptr<OutType>(out),
          out.data_size(),
          static_cast<CTYPE>(start_),
@@ -172,7 +172,6 @@ void ArgReduce::eval_gpu(const std::vector<array>& inputs, array& out) {
          kernel,
          num_blocks,
          block_dim(),
-          0,
          gpu_ptr<T>(in),
          gpu_ptr<uint32_t>(out),
          out.size(),
@@ -16,8 +16,14 @@ namespace cu {

 namespace cg = cooperative_groups;

+constexpr int BINARY_MAX_BLOCK_DIM = 1024;
+
 template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
-__global__ void binary_ss(const In* a, const In* b, Out* out, IdxT size) {
+__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_ss(
+    const In* a,
+    const In* b,
+    Out* out,
+    IdxT size) {
  IdxT index = cg::this_grid().thread_rank();

  if ((index + 1) * N_READS > size) {
@@ -36,7 +42,11 @@ __global__ void binary_ss(const In* a, const In* b, Out* out, IdxT size) {
 }

 template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
-__global__ void binary_sv(const In* a, const In* b, Out* out, IdxT size) {
+__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_sv(
+    const In* a,
+    const In* b,
+    Out* out,
+    IdxT size) {
  IdxT index = cg::this_grid().thread_rank();

  if ((index + 1) * N_READS > size) {
@@ -57,7 +67,11 @@ __global__ void binary_sv(const In* a, const In* b, Out* out, IdxT size) {
 }

 template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
-__global__ void binary_vs(const In* a, const In* b, Out* out, IdxT size) {
+__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_vs(
+    const In* a,
+    const In* b,
+    Out* out,
+    IdxT size) {
  IdxT index = cg::this_grid().thread_rank();

  if ((index + 1) * N_READS > size) {
@@ -78,7 +92,11 @@ __global__ void binary_vs(const In* a, const In* b, Out* out, IdxT size) {
 }

 template <typename Op, typename In, typename Out, typename IdxT, int N_READS>
-__global__ void binary_vv(const In* a, const In* b, Out* out, IdxT size) {
+__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_vv(
+    const In* a,
+    const In* b,
+    Out* out,
+    IdxT size) {
  IdxT index = cg::this_grid().thread_rank();

  if ((index + 1) * N_READS > size) {
@@ -291,7 +309,6 @@ void binary_op_gpu_inplace(
                        kernel,
                        {num_blocks_x, num_blocks_y},
                        block_dims,
-                        0,
                        gpu_ptr<InType>(a),
                        gpu_ptr<InType>(b),
                        gpu_ptr<OutType>(out),
@@ -309,7 +326,6 @@ void binary_op_gpu_inplace(
                      kernel,
                      {num_blocks_x, num_blocks_y},
                      block_dims,
-                      0,
                      gpu_ptr<InType>(a),
                      gpu_ptr<InType>(b),
                      gpu_ptr<OutType>(out),
@@ -333,12 +349,16 @@ void binary_op_gpu_inplace(
              kernel = cu::binary_vv<Op, InType, OutType, IdxT, N_READS>;
            }
            auto [num_blocks, block_dims] = get_launch_args(
-                out.data_size(), out.shape(), out.strides(), large(), N_READS);
+                out.data_size(),
+                out.shape(),
+                out.strides(),
+                large(),
+                N_READS,
+                cu::BINARY_MAX_BLOCK_DIM);
            encoder.add_kernel_node(
                kernel,
                num_blocks,
                block_dims,
-                0,
                gpu_ptr<InType>(a),
                gpu_ptr<InType>(b),
                gpu_ptr<OutType>(out),
@@ -314,7 +314,6 @@ void binary_two_op_gpu_inplace(
                        kernel,
                        {num_blocks_x, num_blocks_y},
                        block_dims,
-                        0,
                        gpu_ptr<InType>(a),
                        gpu_ptr<InType>(b),
                        gpu_ptr<OutType>(out_a),
@@ -333,7 +332,6 @@ void binary_two_op_gpu_inplace(
                      kernel,
                      {num_blocks_x, num_blocks_y},
                      block_dims,
-                      0,
                      gpu_ptr<InType>(a),
                      gpu_ptr<InType>(b),
                      gpu_ptr<OutType>(out_a),
@@ -367,7 +365,6 @@ void binary_two_op_gpu_inplace(
                kernel,
                num_blocks,
                block_dims,
-                0,
                gpu_ptr<InType>(a),
                gpu_ptr<InType>(b),
                gpu_ptr<OutType>(out_a),
@@ -351,7 +351,8 @@ void Compiled::eval_gpu(
  auto [kernel, max_block_dims] = mod.get_kernel_and_dims(kernel_name);
  auto [num_blocks, block_dims] =
      get_launch_args(outputs[0], large, work_per_thread, max_block_dims);
-  encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
+  encoder.add_kernel_node_raw(
+      kernel, num_blocks, block_dims, {}, 0, args.args());
 }

 } // namespace mlx::core
@@ -39,7 +39,7 @@ struct ConvCacheKey {
 };

 auto& conv_cache() {
-  static LRUBytesKeyCache<
+  static thread_local LRUBytesKeyCache<
      ConvCacheKey,
      std::pair<ConvBackendType, std::optional<DnnGraph>>>
      cache("MLX_CUDA_CONV_CACHE_SIZE", /* default_capacity */ 128);
@@ -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().get_cudnn_handle(), dtype, compute_dtype);
+  DnnGraph graph(get_cudnn_handle(encoder.device()), dtype, compute_dtype);
  auto x_ = graph.tensor_nchw("X", 'x', x);
  auto w_ = graph.tensor_nchw("W", 'w', w);

@@ -139,7 +139,7 @@ std::optional<DnnGraph> build_conv_graph(
  if (dtype == float32 && !env::enable_tf32()) {
    graph.deselect_numeric_notes({fe::NumericalNote_t::TENSOR_CORE});
  }
-  CHECK_CUDNN_FE_ERROR(graph.build());
+  CHECK_CUDNN_ERROR(graph.build());
  return graph;
 }

@@ -252,6 +252,10 @@ void register_args(

 } // namespace

+void init_cudnn_conv_cache() {
+  conv_cache();
+}
+
 void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
  nvtx3::scoped_range r("Convolution::eval_gpu");
  if (out_.size() == 0) {
@@ -269,20 +273,19 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {

  // Search cache.
  BytesKey<ConvCacheKey> cache_key;
-  cache_key.pod = {
-      encoder.device().cuda_device(),
-      dtype_to_cudnn_type(dtype),
-      vector_key(in.shape()),
-      vector_key(wt.shape()),
-      vector_key(kernel_strides_),
-      vector_key(padding_lo_),
-      vector_key(padding_hi_),
-      vector_key(kernel_dilation_),
-      groups_,
-      flip_,
-      get_alignment(in),
-      get_alignment(wt),
-      get_alignment(out)};
+  cache_key.pod.device_id = encoder.device().cuda_device();
+  cache_key.pod.cudnn_dtype = dtype_to_cudnn_type(dtype);
+  cache_key.pod.input_shape = vector_key(in.shape());
+  cache_key.pod.weight_shape = vector_key(wt.shape());
+  cache_key.pod.stride = vector_key(kernel_strides_);
+  cache_key.pod.padding_lo = vector_key(padding_lo_);
+  cache_key.pod.padding_hi = vector_key(padding_hi_);
+  cache_key.pod.dilation = vector_key(kernel_dilation_);
+  cache_key.pod.groups = groups_;
+  cache_key.pod.flip = flip_;
+  cache_key.pod.input_alignment = get_alignment(in);
+  cache_key.pod.weight_alignment = get_alignment(wt);
+  cache_key.pod.output_alignment = get_alignment(out);
  if (auto it = conv_cache().find(cache_key); it != conv_cache().end()) {
    auto& [backend_type, graph] = it->second;
    if (graph) {
@@ -290,7 +293,7 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
      std::tie(in, wt, out) =
          prepare_args(encoder, backend_type, in, wt, out, groups_, s);
      register_args(encoder, backend_type, in, wt, out, out_);
-      CHECK_CUDNN_FE_ERROR(graph->encode_capturing(
+      CHECK_CUDNN_ERROR(graph->encode_capturing(
          encoder,
          {
              {'x', gpu_ptr<void>(in)},
@@ -372,7 +375,7 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {

  if (graph) {
    register_args(encoder, backend_type, in, wt, out, out_);
-    CHECK_CUDNN_FE_ERROR(graph->encode_capturing(
+    CHECK_CUDNN_ERROR(graph->encode_capturing(
        encoder,
        {
            {'x', gpu_ptr<void>(in)},
@@ -117,7 +117,6 @@ array unfold_inputs_nd(
        cu::naive_unfold_nd<DataType, NDIM>,
        num_blocks,
        block_dims,
-        0,
        gpu_ptr<DataType>(in),
        gpu_ptr<DataType>(unfolded),
        filter_size,
@@ -120,7 +120,6 @@ array grouped_unfold_transpose_inputs_nd(
        cu::naive_grouped_unfold_transpose_nd<DataType, NDIM>,
        num_blocks,
        block_dims,
-        0,
        gpu_ptr<DataType>(in),
        gpu_ptr<DataType>(unfolded),
        filter_size,
@@ -76,7 +76,6 @@ void copy_contiguous(
            kernel,
            num_blocks,
            block_dims,
-            0,
            gpu_ptr<InType>(in) + in_offset,
            gpu_ptr<OutType>(out) + out_offset,
            out.data_size());
@@ -137,7 +137,6 @@ void copy_general(
                    kernel,
                    {num_blocks_x, num_blocks_y},
                    block_dims,
-                    0,
                    in_ptr,
                    out_ptr,
                    rest,
@@ -154,7 +153,6 @@ void copy_general(
                  kernel,
                  {num_blocks_x, num_blocks_y},
                  block_dims,
-                  0,
                  in_ptr,
                  out_ptr,
                  rest,
@@ -83,7 +83,6 @@ void copy_general_dynamic(
                        dims_constant()>,
                    num_blocks,
                    block_dims,
-                    0,
                    in_ptr,
                    out_ptr,
                    out.size(),
@@ -99,7 +98,6 @@ void copy_general_dynamic(
                  cu::copy_gg_dynamic<InType, OutType, IdxT>,
                  num_blocks,
                  block_dims,
-                  0,
                  in_ptr,
                  out_ptr,
                  out.size(),
@@ -154,7 +154,6 @@ void copy_general_input(
            kernel,
            {num_blocks_x, num_blocks_y},
            block_dims,
-            0,
            in_ptr,
            out_ptr,
            int64_t(shape[0]),
@@ -195,7 +194,6 @@ void copy_general_input(
                    kernel,
                    {num_blocks_x, num_blocks_y},
                    block_dims,
-                    0,
                    in_ptr,
                    out_ptr,
                    rest,
@@ -213,7 +211,6 @@ void copy_general_input(
                  kernel,
                  {num_blocks_x, num_blocks_y},
                  block_dims,
-                  0,
                  in_ptr,
                  out_ptr,
                  rest,
@@ -2,44 +2,13 @@

 #include "mlx/backend/cuda/cublas_utils.h"
 #include "mlx/backend/cuda/cuda.h"
+#include "mlx/backend/gpu/device_info.h"
 #include "mlx/utils.h"

 namespace mlx::core {
+
 namespace cublas_utils {

-namespace {
-
-struct CublasPreference {
-  CublasPreference(cu::Device& device) {
-    // The recommended cublas workspace size is 4 MiB for pre-Hopper and 32 MiB
-    // for Hopper+:
-    // https://docs.nvidia.com/cuda/cublas/#cublassetworkspace
-    uint64_t MiB = 1024 * 1024;
-    uint64_t workspace_size =
-        device.compute_capability_major() >= 9 ? 32 * MiB : 4 * MiB;
-
-    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceCreate(&pref_));
-    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceSetAttribute(
-        pref_,
-        CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
-        &workspace_size,
-        sizeof(uint64_t)));
-  }
-
-  ~CublasPreference() {
-    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceDestroy(pref_));
-  }
-
-  cublasLtMatmulPreference_t pref_{nullptr};
-};
-
-} // namespace
-
-cublasLtMatmulPreference_t get_preference(cu::Device& device) {
-  static CublasPreference pref(device);
-  return pref.pref_;
-}
-
 cublasLtMatrixLayout_t create_matrix_layout(
    cudaDataType_t type,
    uint64_t rows,
@@ -70,6 +39,59 @@ cublasLtMatrixLayout_t create_matrix_layout(

 } // namespace cublas_utils

+namespace {
+
+auto& cublas_handles_cache() {
+  struct CublasHandles {
+    ~CublasHandles() {
+      if (handle) {
+        CHECK_CUBLAS_ERROR(cublasLtDestroy(handle));
+        CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceDestroy(pref));
+      }
+    }
+    cublasLtHandle_t handle{nullptr};
+    cublasLtMatmulPreference_t pref{nullptr};
+  };
+  static thread_local std::vector<CublasHandles> cache(gpu::device_count());
+  return cache;
+}
+
+auto get_cublas_handles(cu::Device& device) {
+  auto& storage = cublas_handles_cache().at(device.cuda_device());
+  if (!storage.handle) {
+    // Create cublasLt handle.
+    device.make_current();
+    CHECK_CUBLAS_ERROR(cublasLtCreate(&storage.handle));
+    // The recommended cublas workspace size is 4 MiB for pre-Hopper and 32
+    // MiB for Hopper+:
+    // https://docs.nvidia.com/cuda/cublas/#cublassetworkspace
+    uint64_t MiB = 1024 * 1024;
+    uint64_t workspace_size =
+        device.compute_capability_major() >= 9 ? 32 * MiB : 4 * MiB;
+    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceCreate(&storage.pref));
+    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceSetAttribute(
+        storage.pref,
+        CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
+        &workspace_size,
+        sizeof(uint64_t)));
+  }
+  return std::make_tuple(storage.handle, storage.pref);
+}
+
+} // namespace
+
+void check_cublas_error(const char* name, cublasStatus_t err) {
+  if (err != CUBLAS_STATUS_SUCCESS) {
+    // TODO: Use cublasGetStatusString when it is widely available.
+    throw std::runtime_error(
+        fmt::format("{} failed with code: {}.", name, static_cast<int>(err)));
+  }
+}
+
+void init_cublas_handles_cache() {
+  cublas_handles_cache();
+}
+
 CublasMatmulBase::~CublasMatmulBase() {
  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(a_desc_));
  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(b_desc_));
@@ -98,20 +120,12 @@ void CublasMatmulBase::init_base(
  M_ = a_rows;
  N_ = b_cols;
  scale_type_ = scale_type;
-  handle_ = device.get_cublaslt_handle();
-  pref_ = cublas_utils::get_preference(device);
+  std::tie(handle_, pref_) = get_cublas_handles(device);
  heuristic_.state = CUBLAS_STATUS_NOT_INITIALIZED;

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

-  int32_t pointer_mode = CUBLASLT_POINTER_MODE_HOST;
-  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
-      matmul_desc_,
-      CUBLASLT_MATMUL_DESC_POINTER_MODE,
-      &pointer_mode,
-      sizeof(int32_t)));
-
  // In cublasLt matrices use column-major layout, while it is possible to use
  // the CUBLASLT_ORDER_ROW option to switch to row-major layout, the bias
  // epilogue does not work with the option. So instead we swap A and B to make
@@ -1,17 +1,15 @@
 // Copyright © 2025 Apple Inc.
 #pragma once

-#include <cublasLt.h>
 #include "mlx/array.h"
 #include "mlx/backend/cuda/device.h"
 #include "mlx/dtype_utils.h"

+#include <cublasLt.h>
+
 namespace mlx::core {
 namespace cublas_utils {

-// Get the shared cublas preference for a device
-cublasLtMatmulPreference_t get_preference(cu::Device& device);
-
 cublasLtMatrixLayout_t create_matrix_layout(
    cudaDataType_t type,
    uint64_t rows,
@@ -42,6 +40,12 @@ inline cudaDataType_t dtype_to_cublas_type(Dtype dtype, std::string_view tag) {

 } // namespace cublas_utils

+void check_cublas_error(const char* name, cublasStatus_t err);
+
+#define CHECK_CUBLAS_ERROR(cmd) check_cublas_error(#cmd, (cmd))
+
+void init_cublas_handles_cache();
+
 class CublasMatmulBase {
 public:
  virtual ~CublasMatmulBase();
@@ -2,23 +2,17 @@

 #pragma once

-#include <cublasLt.h>
 #include <cuda.h>
 #include <cuda_runtime.h>
-#include <cudnn.h>

 namespace mlx::core {

 // Throw exception if the cuda API does not succeed.
-void check_cublas_error(const char* name, cublasStatus_t err);
 void check_cuda_error(const char* name, cudaError_t err);
 void check_cuda_error(const char* name, CUresult err);
-void check_cudnn_error(const char* name, cudnnStatus_t err);

 // The macro version that prints the command that failed.
-#define CHECK_CUBLAS_ERROR(cmd) check_cublas_error(#cmd, (cmd))
 #define CHECK_CUDA_ERROR(cmd) check_cuda_error(#cmd, (cmd))
-#define CHECK_CUDNN_ERROR(cmd) check_cudnn_error(#cmd, (cmd))

 // Base class for RAII managed CUDA resources.
 template <typename Handle, cudaError_t (*Destroy)(Handle)>
@@ -83,6 +77,7 @@ class CudaGraphExec : public CudaHandle<cudaGraphExec_t, cudaGraphExecDestroy> {

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

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

 #include "mlx/backend/cuda/cudnn_utils.h"
 #include "mlx/backend/cuda/device.h"
+#include "mlx/backend/gpu/device_info.h"

 namespace mlx::core {

@@ -47,8 +48,48 @@ inline auto nhwc_to_nchw(const array& x) {
  return std::make_tuple(std::move(shape), std::move(strides));
 }

+auto& cudnn_handles_cache() {
+  struct CudnnHandle {
+    ~CudnnHandle() {
+      if (handle) {
+        CHECK_CUDNN_ERROR(cudnnDestroy(handle));
+      }
+    }
+    cudnnHandle_t handle{nullptr};
+  };
+  static thread_local std::vector<CudnnHandle> cache(gpu::device_count());
+  return cache;
+}
+
 } // namespace

+void check_cudnn_error(const char* name, cudnnStatus_t err) {
+  if (err != CUDNN_STATUS_SUCCESS) {
+    throw std::runtime_error(
+        fmt::format("{} failed: {}.", name, cudnnGetErrorString(err)));
+  }
+}
+
+void check_cudnn_error(const char* name, fe::error_t err) {
+  if (!err.is_good()) {
+    throw std::runtime_error(
+        fmt::format("{} failed: {}.", name, err.get_message()));
+  }
+}
+
+cudnnHandle_t get_cudnn_handle(cu::Device& device) {
+  auto& storage = cudnn_handles_cache().at(device.cuda_device());
+  if (!storage.handle) {
+    device.make_current();
+    CHECK_CUDNN_ERROR(cudnnCreate(&storage.handle));
+  }
+  return storage.handle;
+}
+
+void init_cudnn_handles_cache() {
+  cudnn_handles_cache();
+}
+
 fe::error_t DnnGraph::prepare() {
  RETURN_IF_ERROR(validate());
  try {
@@ -71,10 +112,26 @@ fe::error_t DnnGraph::encode_graph(
    cu::CommandEncoder& encoder,
    std::unordered_map<int64_t, void*> variant_pack) {
  cudnnSetStream(handle_, encoder.stream());
-  CudaGraph cuda_graph(encoder.device());
-  RETURN_IF_ERROR(populate_cuda_graph(
-      handle_, variant_pack, prepare_workspace(encoder), cuda_graph));
-  encoder.add_graph_node(cuda_graph);
+  auto* workspace_ptr = prepare_workspace(encoder);
+  if (!cached_cuda_graph_) {
+    // First call: populate the CUDA graph from the cuDNN execution plan.
+    // Also compute and cache the subgraph key to avoid calling
+    // cudaGraphKernelNodeGetAttribute on every subsequent call (expensive
+    // on WDDM where each driver API call has ~40-400us overhead).
+    cached_cuda_graph_.emplace(encoder.device());
+    RETURN_IF_ERROR(populate_cuda_graph(
+        handle_, variant_pack, workspace_ptr, *cached_cuda_graph_));
+    std::tie(cached_subgraph_key_, cached_is_updatable_) =
+        cu::subgraph_to_key(*cached_cuda_graph_);
+  } else {
+    // Subsequent calls: patch data pointers without re-running kernel setup.
+    RETURN_IF_ERROR(update_cuda_graph(
+        handle_, variant_pack, workspace_ptr, *cached_cuda_graph_));
+  }
+  // Add the cuDNN child graph to the parent CUDA graph for batched launch.
+  // The pre-computed subgraph key avoids expensive per-node attribute queries.
+  encoder.add_graph_node(
+      *cached_cuda_graph_, cached_subgraph_key_, cached_is_updatable_);
  return {};
 }

@@ -93,7 +150,7 @@ fe::error_t DnnGraph::encode_capturing(

 void* DnnGraph::prepare_workspace(cu::CommandEncoder& encoder) {
  int64_t workspace_size = 0;
-  CHECK_CUDNN_FE_ERROR(get_workspace_size(workspace_size));
+  CHECK_CUDNN_ERROR(get_workspace_size(workspace_size));
  return allocate_workspace(encoder, workspace_size);
 }

@@ -2,6 +2,10 @@

 #pragma once

+#include <cassert>
+#include <optional>
+
+#include "mlx/backend/cuda/cuda_utils.h"
 #include "mlx/backend/cuda/device/config.h"
 #include "mlx/backend/cuda/utils.h"
 #include "mlx/dtype_utils.h"
@@ -17,14 +21,16 @@ class CommandEncoder;

 namespace fe = cudnn_frontend;

-#define CHECK_CUDNN_FE_ERROR(cmd)                                    \
-  do {                                                               \
-    auto error = cmd;                                                \
-    if (!error.is_good()) {                                          \
-      throw std::runtime_error(                                      \
-          fmt::format("{} failed: {}.", #cmd, error.get_message())); \
-    }                                                                \
-  } while (0)
+void check_cudnn_error(const char* name, cudnnStatus_t err);
+void check_cudnn_error(const char* name, fe::error_t err);
+
+#define CHECK_CUDNN_ERROR(cmd) check_cudnn_error(#cmd, (cmd))
+
+cudnnHandle_t get_cudnn_handle(cu::Device& device);
+
+void init_cudnn_handles_cache();
+void init_cudnn_conv_cache();
+void init_cudnn_sdpa_cache();

 // Return pointer alignment of |x|'s data.
 inline uint8_t get_alignment(const array& x) {
@@ -123,6 +129,20 @@ class DnnGraph : public fe::graph::Graph {
    return attrs;
  }

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

  cudnnHandle_t handle_;
+  std::optional<CudaGraph> cached_cuda_graph_;
+  std::string cached_subgraph_key_;
+  bool cached_is_updatable_{true};
 };

 } // namespace mlx::core
@@ -373,7 +373,8 @@ void CustomKernel::eval_gpu(
              kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem);
        }
      });
-  encoder.add_kernel_node(kernel, grid, block, shared_memory_, args.args());
+  encoder.add_kernel_node_raw(
+      kernel, grid, block, {}, shared_memory_, args.args());
 }

 } // namespace mlx::core::fast
@@ -22,7 +22,7 @@ inline void check_cutlass_error(const char* name, cutlass::Status status) {
 }

 // The macro version that prints the command that failed.
-#define CHECK_CUTLASS_ERROR(cmd) check_cutlass_error(#cmd, (cmd))
+#define CHECK_CUTLASS_ERROR(cmd) ::mlx::core::check_cutlass_error(#cmd, (cmd))

 // Maps CPU types to CUTLASS types.
 template <typename T>
@@ -1,7 +1,6 @@
 // Copyright © 2025 Apple Inc.

 #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"
@@ -31,6 +30,11 @@ const char* save_cuda_graphs_dot_file() {
  return filename;
 }

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

 Device::Device(int device) : device_(device) {
@@ -42,51 +46,28 @@ Device::Device(int device) : device_(device) {
      &concurrent_managed_access_,
      cudaDevAttrConcurrentManagedAccess,
      device_));
+  CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
+      &host_native_atomic_, cudaDevAttrHostNativeAtomicSupported, device_));
+  CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
+      &managed_memory_, cudaDevAttrManagedMemory, device_));
+  CHECK_CUDA_ERROR(cudaDeviceGetAttribute(
+      &memory_pools_, cudaDevAttrMemoryPoolsSupported, device_));
 }

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

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

-CommandEncoder& Device::get_command_encoder(Stream s) {
-  auto it = encoders_.find(s.index);
-  if (it == encoders_.end()) {
-    it = encoders_.try_emplace(s.index, *this).first;
-  }
-  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()) {
@@ -208,13 +189,7 @@ std::pair<int, int> get_graph_limits(Device& d) {
      mb = 400;
      break;
    case 900: // H100
-      ops = 30;
-      mb = 400;
-      break;
    case 1000: // B200
-      ops = 50;
-      mb = 500;
-      break;
    case 1200: // Consumer Blackwell
      ops = 100;
      mb = 1000;
@@ -231,13 +206,19 @@ CommandEncoder::CommandEncoder(Device& d)
    : device_(d),
      stream_(d),
      graph_(d),
-      worker_(d),
+      worker_(std::make_shared<Worker>(d)),
      graph_cache_("MLX_CUDA_GRAPH_CACHE_SIZE", /* default_capacity */ 400) {
  std::tie(max_ops_per_graph_, max_mb_per_graph_) = get_graph_limits(d);
+  worker_->start();
+}
+
+CommandEncoder::~CommandEncoder() {
+  synchronize();
+  worker_->stop();
 }

 void CommandEncoder::add_completed_handler(std::function<void()> task) {
-  worker_.add_task(std::move(task));
+  worker_->add_task(std::move(task));
 }

 void CommandEncoder::set_input_array(const array& arr) {
@@ -259,51 +240,88 @@ void CommandEncoder::set_output_array(const array& arr) {
  active_outputs_.push_back(id);
 }

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

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

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

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

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

 std::pair<std::string, bool> subgraph_to_key(cudaGraph_t graph) {
@@ -407,6 +437,24 @@ void CommandEncoder::add_graph_node(cudaGraph_t child) {
  insert_graph_dependencies(GraphNode{node, sub_graph_key});
 }

+void CommandEncoder::add_graph_node(
+    cudaGraph_t child,
+    const std::string& subgraph_key,
+    bool is_updatable) {
+  if (!use_cuda_graphs()) {
+    node_count_++;
+    CudaGraphExec graph_exec;
+    graph_exec.instantiate(child);
+    device_.make_current();
+    CHECK_CUDA_ERROR(cudaGraphLaunch(graph_exec, stream()));
+    return;
+  }
+  is_graph_updatable_ &= is_updatable;
+  cudaGraphNode_t node;
+  CHECK_CUDA_ERROR(cudaGraphAddChildGraphNode(&node, graph_, NULL, 0, child));
+  insert_graph_dependencies(GraphNode{node, subgraph_key});
+}
+
 bool CommandEncoder::needs_commit() {
  return (node_count_ > max_ops_per_graph_) ||
      ((bytes_in_graph_ >> 20) > max_mb_per_graph_);
@@ -482,7 +530,7 @@ void CommandEncoder::commit() {
  }

  // Put completion handlers in a batch.
-  worker_.commit(stream_);
+  worker_->commit(stream_);
  node_count_ = 0;
  bytes_in_graph_ = 0;
 }
@@ -497,18 +545,17 @@ void CommandEncoder::synchronize() {
 }

 Device& device(int cuda_device) {
-  static auto devices = []() {
-    std::vector<Device> devices;
+  // The devices are leak intentionally as user code may still be accessing
+  // device after main thread teardown.
+  static auto* devices = []() {
+    auto* devices = new std::vector<Device>;
    int device_count = gpu::device_count();
    for (int i = 0; i < device_count; ++i) {
-      devices.emplace_back(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);
+  return devices->at(cuda_device);
 }

 Device& device(mlx::core::Device d) {
@@ -516,7 +563,18 @@ Device& device(mlx::core::Device d) {
 }

 CommandEncoder& get_command_encoder(Stream s) {
-  return device(s.device).get_command_encoder(s);
+  auto& encoders = get_command_encoders();
+  auto it = encoders.find(s.index);
+  if (it == encoders.end()) {
+    throw std::runtime_error(
+        fmt::format("There is no Stream(gpu, {}) in current thread.", s.index));
+  }
+  return it->second;
+}
+
+std::unordered_map<int, CommandEncoder>& get_command_encoders() {
+  static thread_local std::unordered_map<int, CommandEncoder> encoders;
+  return encoders;
 }

 } // namespace mlx::core::cu
@@ -5,17 +5,19 @@
 #include "mlx/array.h"
 #include "mlx/backend/cuda/allocator.h"
 #include "mlx/backend/cuda/lru_cache.h"
-#include "mlx/backend/cuda/worker.h"
+#include "mlx/backend/cuda/utils.h"
 #include "mlx/stream.h"

-#include <cublasLt.h>
-#include <cuda.h>
-#include <cudnn.h>
-
+#include <memory>
 #include <unordered_map>

 namespace mlx::core::cu {

+// Compute a key and updatability flag for a CUDA graph by walking its nodes.
+std::pair<std::string, bool> subgraph_to_key(cudaGraph_t graph);
+
+class Worker;
+
 class CommandEncoder {
 public:
  struct CaptureContext {
@@ -32,6 +34,7 @@ class CommandEncoder {
  };

  explicit CommandEncoder(Device& d);
+  ~CommandEncoder();

  CommandEncoder(const CommandEncoder&) = delete;
  CommandEncoder& operator=(const CommandEncoder&) = delete;
@@ -47,10 +50,17 @@ class CommandEncoder {
  void set_output_array(const array& arr);

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

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

  void add_graph_node(cudaGraph_t child);
+  void add_graph_node(
+      cudaGraph_t child,
+      const std::string& subgraph_key,
+      bool is_updatable);

  void add_temporary(const array& arr) {
    temporaries_.push_back(arr.data_shared_ptr());
@@ -98,8 +120,8 @@ class CommandEncoder {
  void synchronize();

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

  struct GraphNode {
    cudaGraphNode_t node;
@@ -117,7 +139,7 @@ class CommandEncoder {
  Device& device_;
  CudaStream stream_;
  CudaGraph graph_;
-  Worker worker_;
+  std::shared_ptr<Worker> worker_;
  int node_count_{0};
  bool in_concurrent_{false};
  std::vector<cudaGraphNode_t> from_nodes_;
@@ -148,10 +170,6 @@ class Device {
  // Make this device the current cuda device, this method is thread-safe.
  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_;
  }
@@ -164,20 +182,31 @@ class Device {
  bool concurrent_managed_access() const {
    return concurrent_managed_access_ == 1;
  }
+  bool host_native_atomic() const {
+    return host_native_atomic_ == 1;
+  }
+  bool managed_memory() const {
+    return managed_memory_ == 1;
+  }
+  bool memory_pools() const {
+    return memory_pools_ == 1;
+  }

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

-Device& device(int cuda_device);
-Device& device(mlx::core::Device d);
-CommandEncoder& get_command_encoder(Stream s);
+MLX_API Device& device(int cuda_device);
+MLX_API Device& device(mlx::core::Device d);
+MLX_API CommandEncoder& get_command_encoder(Stream s);
+
+std::unordered_map<int, CommandEncoder>& get_command_encoders();

 } // namespace mlx::core::cu
@@ -0,0 +1,184 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device/utils.cuh"
+
+namespace mlx::core::cu {
+
+__device__ __forceinline__ void hadamard_radix_m(float* x);
+
+template <int N>
+struct Pow2Log2 {
+  static_assert(
+      (N > 0) && ((N & (N - 1)) == 0),
+      "N must be a positive power of two.");
+  static constexpr int value = 1 + Pow2Log2<N / 2>::value;
+};
+
+template <>
+struct Pow2Log2<1> {
+  static constexpr int value = 0;
+};
+
+template <int R>
+__device__ __forceinline__ void hadamard_radix_pow2(float* x) {
+  constexpr int kLogR = Pow2Log2<R>::value;
+  int h = 1;
+#pragma unroll
+  for (int s = 0; s < kLogR; ++s) {
+#pragma unroll
+    for (int i = 0; i < R / 2; ++i) {
+      int k = i & (h - 1);
+      int j = ((i - k) << 1) + k;
+      float a = x[j];
+      float b = x[j + h];
+      x[j] = a + b;
+      x[j + h] = a - b;
+    }
+    h <<= 1;
+  }
+}
+
+template <typename T, int N, int max_radix, int read_width, int stride = 1>
+__global__ void
+hadamard_n(const T* in, T* out, float scale, long long num_transforms) {
+  constexpr int kNumThreads = N / max_radix;
+  constexpr int kLogN = Pow2Log2<N>::value;
+  constexpr int kLogR = Pow2Log2<max_radix>::value;
+  constexpr int kNumSteps = kLogN / kLogR;
+  constexpr int kLogFinal = kLogN % kLogR;
+  constexpr int kFinalRadix = 1 << kLogFinal;
+
+  if (threadIdx.x >= kNumThreads) {
+    return;
+  }
+
+  __shared__ T buf[N];
+  int i = threadIdx.x;
+
+  for (long long transform = blockIdx.x; transform < num_transforms;
+       transform += gridDim.x) {
+    long long base = (transform / stride) * static_cast<long long>(N) * stride +
+        (transform % stride);
+
+    if constexpr (stride == 1) {
+#pragma unroll
+      for (int j = 0; j < max_radix / read_width; ++j) {
+        int index = j * read_width * kNumThreads + i * read_width;
+#pragma unroll
+        for (int r = 0; r < read_width; ++r) {
+          buf[index + r] = in[base + index + r];
+        }
+      }
+    } else {
+#pragma unroll
+      for (int j = 0; j < max_radix; ++j) {
+        buf[j * kNumThreads + i] = in[base + (j * kNumThreads + i) * stride];
+      }
+    }
+    __syncthreads();
+
+    float x[max_radix];
+    int h = 1;
+
+#pragma unroll
+    for (int s = 0; s < kNumSteps; ++s) {
+      int k = i & (h - 1);
+      int j = ((i - k) << kLogR) + k;
+
+#pragma unroll
+      for (int r = 0; r < max_radix; ++r) {
+        x[r] = static_cast<float>(buf[j + h * r]);
+      }
+
+      hadamard_radix_pow2<max_radix>(x);
+
+#pragma unroll
+      for (int r = 0; r < max_radix; ++r) {
+        buf[j + h * r] = static_cast<T>(x[r]);
+      }
+
+      h <<= kLogR;
+      __syncthreads();
+    }
+
+    if constexpr (kFinalRadix > 1) {
+#pragma unroll
+      for (int t = 0; t < max_radix / kFinalRadix; ++t) {
+        int index = i + t * kNumThreads;
+        int k = index & (h - 1);
+        int j = ((index - k) << kLogFinal) + k;
+#pragma unroll
+        for (int r = 0; r < kFinalRadix; ++r) {
+          x[r] = static_cast<float>(buf[j + h * r]);
+        }
+
+        hadamard_radix_pow2<kFinalRadix>(x);
+
+#pragma unroll
+        for (int r = 0; r < kFinalRadix; ++r) {
+          buf[j + h * r] = static_cast<T>(x[r]);
+        }
+      }
+      __syncthreads();
+    }
+
+    if constexpr (stride == 1) {
+#pragma unroll
+      for (int j = 0; j < max_radix / read_width; ++j) {
+        int index = j * read_width * kNumThreads + i * read_width;
+#pragma unroll
+        for (int r = 0; r < read_width; ++r) {
+          float val = static_cast<float>(buf[index + r]);
+          out[base + index + r] = static_cast<T>(val * scale);
+        }
+      }
+    } else {
+#pragma unroll
+      for (int j = 0; j < max_radix; ++j) {
+        out[base + (j * kNumThreads + i) * stride] = buf[j * kNumThreads + i];
+      }
+    }
+
+    __syncthreads();
+  }
+}
+
+template <typename T, int N, int M, int read_width>
+__global__ void
+hadamard_m(const T* in, T* out, float scale, long long num_tasks) {
+  constexpr int kTasksPerBatch = N / read_width;
+
+  for (long long task = blockIdx.x * blockDim.x + threadIdx.x; task < num_tasks;
+       task += blockDim.x * gridDim.x) {
+    long long i = task % kTasksPerBatch;
+    long long batch = task / kTasksPerBatch;
+    long long base = batch * static_cast<long long>(M) * N;
+
+    float x[read_width][M];
+#pragma unroll
+    for (int c = 0; c < M; ++c) {
+#pragma unroll
+      for (int r = 0; r < read_width; ++r) {
+        x[r][c] = static_cast<float>(in[base + c * N + i * read_width + r]);
+      }
+    }
+
+#pragma unroll
+    for (int r = 0; r < read_width; ++r) {
+      hadamard_radix_m(x[r]);
+    }
+
+#pragma unroll
+    for (int c = 0; c < M; ++c) {
+#pragma unroll
+      for (int r = 0; r < read_width; ++r) {
+        out[base + c * N + i * read_width + r] =
+            static_cast<T>(x[r][c] * scale);
+      }
+    }
+  }
+}
+
+} // namespace mlx::core::cu
@@ -65,4 +65,91 @@ __global__ void scatter(
  Op{}(out + out_idx, upd[upd_loc]);
 }

+template <typename T, bool SrcContiguous, bool DstContiguous, typename IdxT>
+__global__ void masked_scatter(
+    const T* dst,
+    const bool* mask,
+    const int32_t* scatter_offsets,
+    const T* src,
+    T* out,
+    IdxT size,
+    IdxT src_batch_size,
+    IdxT mask_batch_size,
+    const __grid_constant__ Shape dst_shape,
+    const __grid_constant__ Strides dst_strides,
+    int32_t dst_ndim,
+    const __grid_constant__ Shape src_shape,
+    const __grid_constant__ Strides src_strides,
+    int32_t src_ndim) {
+  IdxT index = cg::this_grid().thread_rank();
+  if (index >= size) {
+    return;
+  }
+
+  T dst_val;
+  if constexpr (DstContiguous) {
+    dst_val = dst[index];
+  } else {
+    IdxT dst_loc =
+        elem_to_loc(index, dst_shape.data(), dst_strides.data(), dst_ndim);
+    dst_val = dst[dst_loc];
+  }
+
+  if (mask[index]) {
+    IdxT src_index = static_cast<IdxT>(scatter_offsets[index]);
+    if (src_index < src_batch_size) {
+      IdxT batch_idx = index / mask_batch_size;
+      if constexpr (SrcContiguous) {
+        out[index] = src[batch_idx * src_batch_size + src_index];
+      } else {
+        IdxT src_elem = batch_idx * src_batch_size + src_index;
+        IdxT src_loc = elem_to_loc(
+            src_elem, src_shape.data(), src_strides.data(), src_ndim);
+        out[index] = src[src_loc];
+      }
+      return;
+    }
+  }
+
+  out[index] = dst_val;
+}
+
+template <typename T, typename IdxT, int N_READS>
+__global__ void masked_scatter_vec_contiguous(
+    const T* dst,
+    const bool* mask,
+    const int32_t* scatter_offsets,
+    const T* src,
+    T* out,
+    IdxT size,
+    IdxT src_batch_size,
+    IdxT mask_batch_size) {
+  IdxT vec_index = cg::this_grid().thread_rank();
+  IdxT base = vec_index * N_READS;
+  if (base >= size) {
+    return;
+  }
+
+  auto out_vec = load_vector<N_READS>(dst, vec_index, size, static_cast<T>(0));
+  auto mask_vec = load_vector<N_READS>(mask, vec_index, size, false);
+  auto offset_vec = load_vector<N_READS>(scatter_offsets, vec_index, size, 0);
+
+#pragma unroll
+  for (int i = 0; i < N_READS; ++i) {
+    IdxT index = base + i;
+    if (index >= size) {
+      break;
+    }
+    if (mask_vec[i]) {
+      IdxT src_index = static_cast<IdxT>(offset_vec[i]);
+      if (src_index < src_batch_size) {
+        IdxT batch_idx = index / mask_batch_size;
+        out_vec[i] = src[batch_idx * src_batch_size + src_index];
+      }
+    }
+  }
+
+  store_vector<N_READS>(out, vec_index, out_vec, size);
+}
+
 } // namespace mlx::core::cu
@@ -0,0 +1,75 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device/binary_ops.cuh"
+#include "mlx/backend/cuda/device/utils.cuh"
+
+#include <cooperative_groups.h>
+
+namespace mlx::core::cu {
+
+namespace cg = cooperative_groups;
+
+template <
+    typename T,
+    typename IdxT,
+    typename Op,
+    bool OUT_ROW_CONTIG,
+    bool UPD_ROW_CONTIG,
+    bool UPD_SCALAR,
+    int NWORK>
+__global__ void slice_update_op(
+    const T* updates,
+    T* out,
+    int64_t update_size,
+    const __grid_constant__ Shape update_shape,
+    const __grid_constant__ Strides update_strides,
+    int32_t update_ndim,
+    const __grid_constant__ Strides output_strides,
+    int64_t output_offset) {
+  Op op;
+
+  IdxT idx = cg::this_grid().thread_rank() * NWORK;
+  IdxT out_idx;
+  IdxT update_idx;
+
+  if constexpr (OUT_ROW_CONTIG) {
+    out_idx = idx;
+  } else {
+    out_idx = elem_to_loc<IdxT>(
+        idx, update_shape.data(), output_strides.data(), update_ndim);
+  }
+
+  if constexpr (!UPD_SCALAR) {
+    if constexpr (UPD_ROW_CONTIG) {
+      update_idx = idx;
+    } else {
+      update_idx = elem_to_loc<IdxT>(
+          idx, update_shape.data(), update_strides.data(), update_ndim);
+    }
+  } else {
+    update_idx = 0;
+  }
+
+  out += output_offset;
+
+  for (int j = 0; j < NWORK && idx < update_size; j++) {
+    out[out_idx] = op(out[out_idx], updates[update_idx]);
+    idx++;
+
+    if constexpr (OUT_ROW_CONTIG) {
+      out_idx = idx;
+    } else {
+      out_idx += output_strides[update_ndim - 1];
+    }
+
+    if constexpr (UPD_ROW_CONTIG) {
+      update_idx = idx;
+    } else if constexpr (!UPD_SCALAR) {
+      update_idx += update_strides[update_ndim - 1];
+    }
+  }
+}
+
+} // namespace mlx::core::cu
@@ -2,7 +2,9 @@

 #include "mlx/backend/gpu/eval.h"
 #include "mlx/backend/cuda/allocator.h"
-#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/cublas_utils.h"
+#include "mlx/backend/cuda/cudnn_utils.h"
+#include "mlx/backend/cuda/event.h"
 #include "mlx/primitives.h"
 #include "mlx/scheduler.h"

@@ -10,17 +12,32 @@

 namespace mlx::core::gpu {

-void new_stream(Stream s) {
-  // Force initalization of CUDA, so CUDA runtime get destroyed at last.
+void init() {
+  // Force initalization of CUDA, so CUDA runtime get destroyed last.
  cudaFree(nullptr);
  // Make sure CUDA event pool get destroyed after device and stream.
-  cu::CudaEvent::init_pool();
-  // Ensure the static stream objects get created.
-  cu::get_command_encoder(s);
+  mlx::core::cu::CudaEvent::init_pool();
+}
+
+void new_stream(Stream s) {
+  // Make sure the handles get destroyed after CommandEncoder.
+  init_cublas_handles_cache();
+  init_cudnn_handles_cache();
+  init_cudnn_conv_cache();
+  init_cudnn_sdpa_cache();
+  // Create CommandEncoder.
+  assert(s.device == Device::gpu);
+  auto& encoders = cu::get_command_encoders();
+  auto& d = cu::device(s.device);
+  encoders.try_emplace(s.index, d);
 }

 void eval(array& arr) {
  nvtx3::scoped_range r("gpu::eval");
+  // Ensure CUDA context is active on this thread. Required when MLX is called
+  // from threads that have not yet established a CUDA context (e.g. thread
+  // pools, language runtimes that migrate work across OS threads).
+  cu::device(arr.primitive().stream().device).make_current();
  auto outputs = arr.outputs();
  {
    // If the array is a tracer hold a reference
@@ -63,4 +80,8 @@ void synchronize(Stream s) {
  cu::get_command_encoder(s).synchronize();
 }

+void clear_streams() {
+  cu::get_command_encoders().clear();
+}
+
 } // namespace mlx::core::gpu
@@ -174,62 +174,95 @@ class CopyableCudaEvent {
 // AtomicEvent implementations
 ///////////////////////////////////////////////////////////////////////////////

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

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

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

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

-bool supports_concurrent_managed_access() {
-  static bool concurrent_managed_access = []() {
+auto check_gpu_coherency() {
+  static auto coherency = []() {
    int device_count = gpu::device_count();
+    bool concurrent_managed_access = true;
+    bool host_native_atomic = true;
    for (int i = 0; i < device_count; ++i) {
-      if (!cu::device(i).concurrent_managed_access()) {
-        return false;
-      }
+      auto& d = cu::device(i);
+      concurrent_managed_access &= d.concurrent_managed_access();
+      host_native_atomic &= d.host_native_atomic();
    }
-    return true;
+    return std::make_tuple(concurrent_managed_access, host_native_atomic);
  }();
-  return concurrent_managed_access;
+  return coherency;
 }

-AtomicEvent::AtomicEvent() {
-  if (!supports_concurrent_managed_access()) {
-    throw std::runtime_error(
-        "Device does not support synchronization in managed memory.");
+AtomicEvent::AtomicEvent(Device& d) {
+  void* buf;
+  cudaError_t (*cuda_free)(void*);
+  // There are 3 kinds of systems we are implementing for:
+  // 1. concurrentManagedAccess == true
+  //    => use cuda::atom_ref on managed memory
+  // 2. hostNativeAtomicSupported == true
+  //    => use cuda::atom_ref on pinned host memory
+  // 2. no hardware cpu/gpu coherency
+  //    => use cuda::atom_ref on device memory
+  d.make_current();
+  auto [concurrent_managed_access, host_native_atomic] = check_gpu_coherency();
+  if (concurrent_managed_access) {
+    CHECK_CUDA_ERROR(cudaMallocManaged(&buf, sizeof(uint32_t)));
+    cuda_free = cudaFree;
+    coherent_ = true;
+  } else if (host_native_atomic) {
+    CHECK_CUDA_ERROR(cudaMallocHost(&buf, sizeof(uint32_t)));
+    cuda_free = cudaFreeHost;
+    coherent_ = true;
+  } else {
+    CHECK_CUDA_ERROR(cudaMalloc(&buf, sizeof(uint32_t)));
+    cuda_free = cudaFree;
+    coherent_ = false;
+  }
+  buf_ = std::shared_ptr<void>(
+      buf, [cuda_free](void* buf) { CHECK_CUDA_ERROR(cuda_free(buf)); });
+  if (coherent_) {
+    *ptr() = 0;
+  } else {
+    CHECK_CUDA_ERROR(cudaMemset(buf, 0, sizeof(uint32_t)));
  }
-  buf_ = std::shared_ptr<Buffer>(
-      new Buffer{allocator().malloc(sizeof(Atomic))}, [](Buffer* ptr) {
-        allocator().free(*ptr);
-        delete ptr;
-      });
-  *static_cast<uint64_t*>(buf_->raw_ptr()) = 0;
 }

-void AtomicEvent::wait(uint64_t value) {
+void AtomicEvent::wait(uint32_t value) {
  nvtx3::scoped_range r("cu::AtomicEvent::wait");
-  event_wait(atomic(), value);
+  if (coherent_) {
+    event_wait(ptr(), value);
+  } else {
+    while (!is_signaled(value)) {
+      std::this_thread::yield();
+    }
+  }
 }

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

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

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

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

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

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

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

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

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

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

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

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

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

 void Fence::wait(Stream s, const array&) {
@@ -29,15 +30,9 @@ void Fence::update(Stream s, const array& a, bool cross_device) {
    auto& cbuf =
        *static_cast<cu::CudaBuffer*>(const_cast<array&>(a).buffer().ptr());
    if (cbuf.device != -1) {
-      void* new_data;
-      CHECK_CUDA_ERROR(cudaMallocManaged(&new_data, cbuf.size));
-      cbuf.device = -1;
-      auto& encoder = cu::device(s.device).get_command_encoder(s);
+      auto& encoder = cu::get_command_encoder(s);
      encoder.commit();
-      CHECK_CUDA_ERROR(cudaMemcpyAsync(
-          new_data, cbuf.data, cbuf.size, cudaMemcpyDefault, encoder.stream()));
-      CHECK_CUDA_ERROR(cudaFreeAsync(cbuf.data, encoder.stream()));
-      cbuf.data = new_data;
+      cu::allocator().move_to_unified_memory(cbuf, encoder.stream());
    }
  }
  fence->count++;
@@ -0,0 +1,443 @@
+// Copyright © 2025 Apple Inc.
+
+#include <cufftXt.h>
+#include <algorithm>
+#include <cstdint>
+#include <memory>
+#include <numeric>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+#include <cooperative_groups.h>
+#include <nvtx3/nvtx3.hpp>
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cuda/allocator.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/device/complex.cuh"
+#include "mlx/backend/cuda/lru_cache.h"
+#include "mlx/backend/cuda/utils.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename T>
+__global__ void scale_fft_output(T* out, T scale, size_t size) {
+  auto index = cg::this_grid().thread_rank();
+  if (index < size) {
+    out[index] *= scale;
+  }
+}
+
+} // namespace cu
+
+namespace {
+
+void check_cufft_error(const char* name, cufftResult err) {
+  if (err != CUFFT_SUCCESS) {
+    throw std::runtime_error(
+        std::string(name) +
+        " failed with code: " + std::to_string(static_cast<int>(err)) + ".");
+  }
+}
+
+#define CHECK_CUFFT_ERROR(cmd) check_cufft_error(#cmd, (cmd))
+
+enum class FFTTransformType : uint8_t {
+  C2C = 0,
+  R2C = 1,
+  C2R = 2,
+};
+
+struct FFTPlanKey {
+  int device_id;
+  FFTTransformType transform_type;
+  int64_t n;
+  int64_t batch;
+};
+
+struct CuFFTPlan {
+  explicit CuFFTPlan(int device_id, cufftHandle handle, size_t workspace_size)
+      : device_id(device_id), handle(handle), workspace_size(workspace_size) {}
+
+  ~CuFFTPlan() {
+    if (handle != 0) {
+      try {
+        cu::device(device_id).make_current();
+        cufftDestroy(handle);
+      } catch (...) {
+      }
+    }
+  }
+
+  int device_id;
+  cufftHandle handle;
+  size_t workspace_size;
+};
+
+struct OrderedArray {
+  array arr;
+  std::vector<int> order;
+};
+
+auto& fft_plan_cache() {
+  static LRUBytesKeyCache<FFTPlanKey, std::shared_ptr<CuFFTPlan>> cache(
+      "MLX_CUDA_FFT_CACHE_SIZE",
+      /* default_capacity */ 128);
+  return cache;
+}
+
+FFTPlanKey make_plan_key(
+    int device_id,
+    FFTTransformType transform_type,
+    int64_t n,
+    int64_t batch) {
+  FFTPlanKey key{};
+  key.device_id = device_id;
+  key.transform_type = transform_type;
+  key.n = n;
+  key.batch = batch;
+  return key;
+}
+
+cudaDataType_t input_type(FFTTransformType transform_type) {
+  switch (transform_type) {
+    case FFTTransformType::C2C:
+    case FFTTransformType::C2R:
+      return CUDA_C_32F;
+    case FFTTransformType::R2C:
+      return CUDA_R_32F;
+  }
+  throw std::runtime_error("[FFT] Unsupported cuFFT input transform type.");
+}
+
+cudaDataType_t output_type(FFTTransformType transform_type) {
+  switch (transform_type) {
+    case FFTTransformType::C2C:
+    case FFTTransformType::R2C:
+      return CUDA_C_32F;
+    case FFTTransformType::C2R:
+      return CUDA_R_32F;
+  }
+  throw std::runtime_error("[FFT] Unsupported cuFFT output transform type.");
+}
+
+cudaDataType_t execution_type(FFTTransformType transform_type) {
+  switch (transform_type) {
+    case FFTTransformType::C2C:
+      return CUDA_C_32F;
+    case FFTTransformType::R2C:
+      return CUDA_R_32F;
+    case FFTTransformType::C2R:
+      return CUDA_C_32F;
+  }
+  throw std::runtime_error("[FFT] Unsupported cuFFT execution transform type.");
+}
+
+int64_t input_embed(FFTTransformType transform_type, int64_t n) {
+  return transform_type == FFTTransformType::C2R ? (n / 2 + 1) : n;
+}
+
+int64_t output_embed(FFTTransformType transform_type, int64_t n) {
+  return transform_type == FFTTransformType::R2C ? (n / 2 + 1) : n;
+}
+
+int exec_direction(FFTTransformType transform_type, bool inverse) {
+  switch (transform_type) {
+    case FFTTransformType::C2C:
+      return inverse ? CUFFT_INVERSE : CUFFT_FORWARD;
+    case FFTTransformType::R2C:
+      return CUFFT_FORWARD;
+    case FFTTransformType::C2R:
+      return CUFFT_INVERSE;
+  }
+  throw std::runtime_error("[FFT] Unsupported cuFFT execution direction.");
+}
+
+std::shared_ptr<CuFFTPlan> get_fft_plan(
+    cu::CommandEncoder& encoder,
+    FFTTransformType transform_type,
+    int64_t n,
+    int64_t batch) {
+  auto key = BytesKey<FFTPlanKey>{};
+  key.pod =
+      make_plan_key(encoder.device().cuda_device(), transform_type, n, batch);
+
+  auto& cache = fft_plan_cache();
+  if (auto entry = cache.find(key); entry != cache.end()) {
+    return entry->second;
+  }
+
+  encoder.device().make_current();
+
+  cufftHandle handle = 0;
+  size_t workspace_size = 0;
+  try {
+    CHECK_CUFFT_ERROR(cufftCreate(&handle));
+    CHECK_CUFFT_ERROR(cufftSetAutoAllocation(handle, 0));
+    CHECK_CUFFT_ERROR(cufftSetStream(handle, encoder.stream()));
+
+    long long plan_n[1] = {n};
+    long long inembed[1] = {input_embed(transform_type, n)};
+    long long onembed[1] = {output_embed(transform_type, n)};
+    CHECK_CUFFT_ERROR(cufftXtMakePlanMany(
+        handle,
+        /* rank= */ 1,
+        plan_n,
+        inembed,
+        /* istride= */ 1,
+        /* idist= */ input_embed(transform_type, n),
+        input_type(transform_type),
+        onembed,
+        /* ostride= */ 1,
+        /* odist= */ output_embed(transform_type, n),
+        output_type(transform_type),
+        batch,
+        &workspace_size,
+        execution_type(transform_type)));
+  } catch (...) {
+    if (handle != 0) {
+      encoder.device().make_current();
+      cufftDestroy(handle);
+    }
+    throw;
+  }
+
+  auto plan = std::make_shared<CuFFTPlan>(
+      encoder.device().cuda_device(), handle, workspace_size);
+  return cache.emplace(key, plan).first->second;
+}
+
+std::vector<int> make_identity_order(int ndim) {
+  std::vector<int> order(ndim);
+  std::iota(order.begin(), order.end(), 0);
+  return order;
+}
+
+std::vector<int> move_axis_to_back_permutation(int ndim, int axis_pos) {
+  std::vector<int> perm;
+  perm.reserve(ndim);
+  for (int i = 0; i < ndim; ++i) {
+    if (i != axis_pos) {
+      perm.push_back(i);
+    }
+  }
+  perm.push_back(axis_pos);
+  return perm;
+}
+
+std::vector<int> apply_permutation(
+    const std::vector<int>& values,
+    const std::vector<int>& perm) {
+  std::vector<int> out(perm.size());
+  for (int i = 0; i < perm.size(); ++i) {
+    out[i] = values[perm[i]];
+  }
+  return out;
+}
+
+int find_axis_position(const std::vector<int>& order, int axis) {
+  auto it = std::find(order.begin(), order.end(), axis);
+  if (it == order.end()) {
+    throw std::runtime_error("[FFT] Internal axis tracking mismatch.");
+  }
+  return static_cast<int>(it - order.begin());
+}
+
+OrderedArray prepare_input(
+    const OrderedArray& current,
+    int axis,
+    bool allow_direct,
+    cu::CommandEncoder& encoder,
+    Stream s) {
+  int axis_pos = find_axis_position(current.order, axis);
+  bool axis_last = axis_pos == static_cast<int>(current.order.size()) - 1;
+  bool direct = allow_direct && axis_last && current.arr.flags().row_contiguous;
+
+  if (direct) {
+    return current;
+  }
+
+  array view = current.arr;
+  std::vector<int> order = current.order;
+  if (!axis_last) {
+    auto perm = move_axis_to_back_permutation(current.arr.ndim(), axis_pos);
+    view = transpose_in_eval(current.arr, perm);
+    order = apply_permutation(current.order, perm);
+  }
+
+  array packed = contiguous_copy_gpu(view, s);
+  encoder.add_temporary(packed);
+  return {std::move(packed), std::move(order)};
+}
+
+void execute_fft(
+    const array& in,
+    array& out,
+    FFTTransformType transform_type,
+    bool inverse,
+    cu::CommandEncoder& encoder) {
+  if (!in.flags().row_contiguous || in.strides(-1) != 1) {
+    throw std::runtime_error("[FFT] Expected packed row-contiguous FFT input.");
+  }
+
+  int64_t n =
+      transform_type == FFTTransformType::C2R ? out.shape(-1) : in.shape(-1);
+  int64_t batch = in.shape().empty() ? 1 : in.size() / in.shape(-1);
+  auto plan = get_fft_plan(encoder, transform_type, n, batch);
+
+  encoder.set_input_array(in);
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+  encoder.set_output_array(out);
+  encoder.add_completed_handler([plan]() {});
+
+  encoder.device().make_current();
+  CHECK_CUFFT_ERROR(cufftSetStream(plan->handle, encoder.stream()));
+  auto* workspace = allocate_workspace(encoder, plan->workspace_size);
+  CHECK_CUFFT_ERROR(cufftSetWorkArea(plan->handle, workspace));
+
+  auto capture = encoder.capture_context();
+  CHECK_CUFFT_ERROR(cufftXtExec(
+      plan->handle,
+      gpu_ptr<void>(in),
+      gpu_ptr<void>(out),
+      exec_direction(transform_type, inverse)));
+}
+
+void restore_output_layout(const OrderedArray& current, array& out) {
+  Strides out_strides(out.ndim());
+  for (int i = 0; i < current.order.size(); ++i) {
+    out_strides[current.order[i]] = current.arr.strides(i);
+  }
+
+  auto [data_size, row_contiguous, col_contiguous] =
+      check_contiguity(out.shape(), out_strides);
+  bool contiguous =
+      current.arr.flags().contiguous && data_size == current.arr.data_size();
+
+  out.copy_shared_buffer(
+      current.arr,
+      out_strides,
+      {contiguous, row_contiguous, col_contiguous},
+      current.arr.data_size());
+}
+
+void apply_inverse_scale(
+    array& arr,
+    const std::vector<size_t>& axes,
+    const array& out,
+    cu::CommandEncoder& encoder) {
+  if (axes.empty()) {
+    return;
+  }
+
+  double scale = 1.0;
+  for (auto axis : axes) {
+    scale /= out.shape(axis);
+  }
+
+  size_t size = arr.data_size();
+  dim3 block_dims(256);
+  dim3 grid_dims((size + block_dims.x - 1) / block_dims.x);
+
+  encoder.set_input_array(arr);
+  encoder.set_output_array(arr);
+
+  if (arr.dtype() == float32) {
+    float scale_f = static_cast<float>(scale);
+    encoder.add_kernel_node(
+        cu::scale_fft_output<float>,
+        grid_dims,
+        block_dims,
+        gpu_ptr<float>(arr),
+        scale_f,
+        size);
+  } else if (arr.dtype() == complex64) {
+    cu::complex64_t scale_f(static_cast<float>(scale), 0.0f);
+    encoder.add_kernel_node(
+        cu::scale_fft_output<cu::complex64_t>,
+        grid_dims,
+        block_dims,
+        gpu_ptr<cu::complex64_t>(arr),
+        scale_f,
+        size);
+  } else {
+    throw std::runtime_error("[FFT] Unsupported dtype for inverse scaling.");
+  }
+}
+
+} // namespace
+
+void FFT::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("FFT::eval_gpu");
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+  auto& in = inputs[0];
+
+  if (out.size() == 0) {
+    return;
+  }
+
+  auto order = make_identity_order(in.ndim());
+  OrderedArray current{in, std::move(order)};
+
+  std::vector<int> axis_sequence;
+  axis_sequence.reserve(axes_.size());
+  if (inverse_) {
+    for (auto axis : axes_) {
+      axis_sequence.push_back(static_cast<int>(axis));
+    }
+  } else {
+    for (int i = static_cast<int>(axes_.size()) - 1; i >= 0; --i) {
+      axis_sequence.push_back(static_cast<int>(axes_[i]));
+    }
+  }
+
+  int real_axis = axes_.empty() ? -1 : static_cast<int>(axes_.back());
+
+  for (int i = 0; i < axis_sequence.size(); ++i) {
+    int axis = axis_sequence[i];
+    bool step_real = real_ && axis == real_axis;
+    auto transform_type = step_real
+        ? (inverse_ ? FFTTransformType::C2R : FFTTransformType::R2C)
+        : FFTTransformType::C2C;
+
+    // cuFFT may overwrite the input buffer for C2R, so only use the direct
+    // input when the transform is out-of-place from the library's perspective
+    // or when the original input may be donated to the output.
+    auto prepared = prepare_input(
+        current,
+        axis,
+        /* allow_direct= */ transform_type != FFTTransformType::C2R ||
+            is_donatable(in, out),
+        encoder,
+        s);
+
+    Shape step_shape = prepared.arr.shape();
+    if (step_real) {
+      step_shape.back() = out.shape(axis);
+    }
+
+    Dtype step_dtype =
+        transform_type == FFTTransformType::C2R ? float32 : complex64;
+    array step_out(std::move(step_shape), step_dtype, nullptr, {});
+    execute_fft(prepared.arr, step_out, transform_type, inverse_, encoder);
+    encoder.add_temporary(step_out);
+
+    current = {std::move(step_out), std::move(prepared.order)};
+  }
+
+  if (inverse_) {
+    apply_inverse_scale(current.arr, axes_, out, encoder);
+  }
+
+  restore_output_layout(current, out);
+}
+
+} // namespace mlx::core
@@ -0,0 +1,176 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/device/utils.cuh"
+#include "mlx/backend/cuda/gemms/block_mask.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/dtype_utils.h"
+
+#include <cooperative_groups.h>
+
+namespace mlx::core {
+
+namespace cg = cooperative_groups;
+
+namespace cu {
+
+template <typename T, typename MaskT, bool SrcContiguous>
+__global__ void block_mask_copy_kernel(
+    const T* src,
+    T* dst,
+    int block_size,
+    int64_t rows,
+    int64_t cols,
+    const __grid_constant__ Shape src_shape,
+    const __grid_constant__ Strides src_strides,
+    int src_ndim,
+    MaskT* mask,
+    const __grid_constant__ Shape mask_shape,
+    const __grid_constant__ Strides mask_strides,
+    int mask_ndim,
+    int64_t mask_row_stride,
+    int64_t mask_col_stride,
+    int64_t mask_mat_size,
+    int64_t batch_count) {
+  int64_t mat_size = rows * cols;
+  int64_t idx = cg::this_grid().thread_rank();
+  if (idx >= batch_count * mat_size)
+    return;
+
+  int64_t batch = idx / mat_size;
+  int64_t within = idx % mat_size;
+  int64_t mask_batch_offset = elem_to_loc(
+      batch * mask_mat_size, mask_shape.data(), mask_strides.data(), mask_ndim);
+  MaskT mask_val = mask
+      [mask_batch_offset + (within / cols) / block_size * mask_row_stride +
+       (within % cols) / block_size * mask_col_stride];
+
+  int64_t src_offset;
+  if constexpr (SrcContiguous) {
+    src_offset = idx;
+  } else {
+    src_offset = elem_to_loc(
+        batch * mat_size + within,
+        src_shape.data(),
+        src_strides.data(),
+        src_ndim);
+  }
+
+  if constexpr (std::is_same_v<MaskT, bool>) {
+    dst[idx] = mask_val ? src[src_offset] : T(0);
+  } else {
+    dst[idx] = src[src_offset] * T(mask_val);
+  }
+}
+
+} // namespace cu
+
+namespace {
+
+template <typename T, typename F>
+void dispatch_mask_type(Dtype mask_dtype, F&& f) {
+  if (mask_dtype == bool_) {
+    f.template operator()<bool>();
+  } else {
+    f.template operator()<T>();
+  }
+}
+
+void block_mask_copy(
+    cu::CommandEncoder& encoder,
+    const array& src,
+    array& dst,
+    const array& mask,
+    int block_size,
+    int64_t rows,
+    int64_t cols,
+    bool src_contiguous,
+    int64_t batch_count) {
+  int mask_ndim = mask.ndim();
+  int64_t mask_row_str = mask.strides()[mask_ndim - 2];
+  int64_t mask_col_str = mask.strides()[mask_ndim - 1];
+  int64_t mask_mat_size =
+      int64_t(mask.shape()[mask_ndim - 2]) * mask.shape()[mask_ndim - 1];
+
+  auto [num_blocks, block_dims] = get_launch_args(src, src.size() > INT32_MAX);
+
+  dispatch_float_types(src.dtype(), "block_mask_copy", [&](auto type_tag) {
+    using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+
+    dispatch_mask_type<T>(mask.dtype(), [&]<typename MaskT>() {
+      dispatch_bool(src_contiguous, [&](auto contiguous_tag) {
+        constexpr bool Contiguous = decltype(contiguous_tag)::value;
+        encoder.add_kernel_node(
+            cu::block_mask_copy_kernel<T, MaskT, Contiguous>,
+            num_blocks,
+            block_dims,
+            gpu_ptr<T>(src),
+            gpu_ptr<T>(dst),
+            block_size,
+            rows,
+            cols,
+            const_param(src.shape()),
+            const_param(src.strides()),
+            src.ndim(),
+            gpu_ptr<MaskT>(mask),
+            const_param(mask.shape()),
+            const_param(mask.strides()),
+            mask_ndim,
+            mask_row_str,
+            mask_col_str,
+            mask_mat_size,
+            batch_count);
+      });
+    });
+  });
+}
+
+} // namespace
+
+void apply_block_mask(
+    cu::CommandEncoder& encoder,
+    array& data,
+    const array& mask,
+    int block_size,
+    int64_t rows,
+    int64_t cols,
+    int64_t batch_count) {
+  encoder.set_input_array(mask);
+  encoder.set_output_array(data);
+
+  // Use block_mask_copy in-place (src == dst) with SrcContiguous=true.
+  block_mask_copy(
+      encoder, data, data, mask, block_size, rows, cols, true, batch_count);
+}
+
+array copy_with_block_mask(
+    cu::CommandEncoder& encoder,
+    const array& src,
+    const array& mask,
+    int block_size,
+    int64_t rows,
+    int64_t cols,
+    int64_t batch_count) {
+  array dst(src.shape(), src.dtype(), nullptr, {});
+  dst.set_data(cu::malloc_async(dst.nbytes(), encoder));
+  encoder.add_temporary(dst);
+
+  encoder.set_input_array(src);
+  encoder.set_input_array(mask);
+  encoder.set_output_array(dst);
+
+  block_mask_copy(
+      encoder,
+      src,
+      dst,
+      mask,
+      block_size,
+      rows,
+      cols,
+      src.flags().row_contiguous,
+      batch_count);
+
+  return dst;
+}
+
+} // namespace mlx::core
@@ -0,0 +1,28 @@
+// Copyright © 2026 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+#include "mlx/backend/cuda/device.h"
+
+namespace mlx::core {
+
+void apply_block_mask(
+    cu::CommandEncoder& encoder,
+    array& data,
+    const array& mask,
+    int block_size,
+    int64_t rows,
+    int64_t cols,
+    int64_t batch_count);
+
+array copy_with_block_mask(
+    cu::CommandEncoder& encoder,
+    const array& src,
+    const array& mask,
+    int block_size,
+    int64_t rows,
+    int64_t cols,
+    int64_t batch_count);
+
+} // namespace mlx::core
@@ -73,6 +73,14 @@ CublasGemm::CublasGemm(
      batch_count,
      a_batch_stride,
      b_batch_stride);
+
+  // alpha and beta are both host pointers
+  cublasLtPointerMode_t pointer_mode = CUBLASLT_POINTER_MODE_HOST;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_POINTER_MODE,
+      &pointer_mode,
+      sizeof(pointer_mode)));
 }

 CublasGemm::CublasGemm(
@@ -215,8 +223,8 @@ void CublasGemm::execute(
    const void* a,
    const void* b,
    const void* c,
-    float alpha /* = 1 */,
-    float beta /* = 0 */) {
+    const float alpha /* = 1 */,
+    const float beta /* = 0 */) {
  const void* alpha_ptr = &alpha;
  const void* beta_ptr = &beta;
  complex64_t alpha_c, beta_c;
@@ -182,7 +182,6 @@ void CublasGemm::run_batched(
          cu::set_mm_device_pointers_nd<ndim_constant()>,
          num_blocks,
          block_dims,
-          0,
          gpu_ptr<int8_t*>(pointers),
          gpu_ptr<int8_t>(a),
          gpu_ptr<int8_t>(b),
@@ -199,7 +198,6 @@ void CublasGemm::run_batched(
        cu::set_mm_device_pointers_g,
        num_blocks,
        block_dims,
-        0,
        gpu_ptr<int8_t*>(pointers),
        gpu_ptr<int8_t>(a),
        gpu_ptr<int8_t>(b),
@@ -270,7 +268,6 @@ void CublasGemm::run_batched(
          cu::set_addmm_device_pointers_nd<ndim_constant()>,
          num_blocks,
          block_dims,
-          0,
          gpu_ptr<int8_t*>(pointers),
          gpu_ptr<int8_t>(a),
          gpu_ptr<int8_t>(b),
@@ -289,7 +286,6 @@ void CublasGemm::run_batched(
        cu::set_addmm_device_pointers_g,
        num_blocks,
        block_dims,
-        0,
        gpu_ptr<int8_t*>(pointers),
        gpu_ptr<int8_t>(a),
        gpu_ptr<int8_t>(b),
@@ -0,0 +1,339 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/cutlass_utils.cuh"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/dtype_utils.h"
+
+#include <cutlass/epilogue/collective/collective_epilogue.hpp>
+#include <cutlass/epilogue/thread/linear_combination.h>
+#include <cutlass/gemm/collective/collective_mma.hpp>
+#include <cutlass/gemm/device/gemm_universal_adapter.h>
+#include <cutlass/gemm/dispatch_policy.hpp>
+#include <cutlass/gemm/kernel/gemm_universal.hpp>
+
+// We can't put kernel code in mlx::core due to name conflicts of "Shape".
+namespace cutlass_gemm {
+
+using namespace cute;
+
+// Modified from cutlass/include/cutlass/gemm/kernel/sm70_gemm.hpp to fuse
+// gather into GEMM.
+template <
+    class ProblemShape_,
+    class CollectiveMainloop_,
+    class CollectiveEpilogue_>
+class GatherGemm {
+ public:
+  using ProblemShape = ProblemShape_;
+  using CollectiveMainloop = CollectiveMainloop_;
+  using TileShape = typename CollectiveMainloop::TileShape;
+  using TiledMma = typename CollectiveMainloop::TiledMma;
+  using ArchTag = typename CollectiveMainloop::ArchTag;
+  using ElementA = typename CollectiveMainloop::ElementA;
+  using StrideA = typename CollectiveMainloop::StrideA;
+  using ElementB = typename CollectiveMainloop::ElementB;
+  using StrideB = typename CollectiveMainloop::StrideB;
+  using DispatchPolicy = typename CollectiveMainloop::DispatchPolicy;
+  using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator;
+
+  using CollectiveEpilogue = CollectiveEpilogue_;
+  using ElementC = typename CollectiveEpilogue::ElementC;
+  using StrideC = typename CollectiveEpilogue::StrideC;
+  using ElementD = typename CollectiveEpilogue::ElementD;
+  using StrideD = typename CollectiveEpilogue::StrideD;
+
+  static_assert(
+      cute::is_same_v<
+          ElementAccumulator,
+          typename CollectiveEpilogue::ElementAccumulator>,
+      "Mainloop and epilogue do not agree on accumulator value type.");
+
+  static constexpr int SharedStorageSize = static_cast<int>(cute::max(
+      sizeof(typename CollectiveMainloop::SharedStorage),
+      sizeof(typename CollectiveEpilogue::SharedStorage)));
+  static constexpr uint32_t MaxThreadsPerBlock =
+      CUTE_STATIC_V(size(TiledMma{}));
+  static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
+
+  struct Arguments {
+    ProblemShape problem_shape;
+    const uint32_t* lhs_indices;
+    const uint32_t* rhs_indices;
+    typename CollectiveMainloop::Arguments mainloop;
+    typename CollectiveEpilogue::Arguments epilogue;
+  };
+
+  struct Params {
+    ProblemShape problem_shape;
+    const uint32_t* lhs_indices;
+    const uint32_t* rhs_indices;
+    typename CollectiveMainloop::Params mainloop;
+    typename CollectiveEpilogue::Params epilogue;
+  };
+
+  static Params to_underlying_arguments(
+      const Arguments& args,
+      void* workspace) {
+    return {
+        args.problem_shape,
+        args.lhs_indices,
+        args.rhs_indices,
+        CollectiveMainloop::to_underlying_arguments(
+            args.problem_shape, args.mainloop, workspace),
+        CollectiveEpilogue::to_underlying_arguments(
+            args.problem_shape, args.epilogue, workspace)};
+  }
+
+  static cutlass::Status
+  initialize_workspace(const Arguments&, void*, cudaStream_t, void*) {
+    return cutlass::Status::kSuccess;
+  }
+
+  static dim3 get_grid_shape(const Params& params) {
+    auto [m, n, k, l] = params.problem_shape;
+    return dim3{
+        uint32_t(ceil_div(m, shape<0>(TileShape{}))),
+        uint32_t(ceil_div(n, shape<1>(TileShape{}))),
+        uint32_t(l)};
+  }
+
+  static dim3 get_block_shape() {
+    return dim3{MaxThreadsPerBlock, 1, 1};
+  }
+
+  CUTLASS_DEVICE void operator()(const Params& params, char* smem_buf) {
+    int thread_idx = int(threadIdx.x);
+    int m_coord = int(blockIdx.x);
+    int n_coord = int(blockIdx.y);
+    int l_coord = int(blockIdx.z);
+
+    auto shape_MNKL = append<4>(params.problem_shape, Int<1>{});
+    auto cta_tile = TileShape{};
+    auto cta_coord = make_coord(m_coord, n_coord, _, l_coord);
+
+    // Represent the full tensors.
+    Tensor mA_mkl = make_tensor(
+        make_gmem_ptr(params.mainloop.ptr_A),
+        select<0, 2, 3>(shape_MNKL),
+        params.mainloop.dA);
+    Tensor mB_nkl = make_tensor(
+        make_gmem_ptr(params.mainloop.ptr_B),
+        select<1, 2, 3>(shape_MNKL),
+        params.mainloop.dB);
+
+    // Get batch slice.
+    Tensor mA_mk = mA_mkl(_, _, params.lhs_indices[l_coord]);
+    Tensor mB_nk = mB_nkl(_, _, params.rhs_indices[l_coord]);
+
+    // Slice to get the tiles this thread block is responsible for.
+    Tensor gA =
+        local_tile(mA_mk, cta_tile, take<0, 3>(cta_coord), Step<_1, X, _1>{});
+    Tensor gB =
+        local_tile(mB_nk, cta_tile, take<0, 3>(cta_coord), Step<X, _1, _1>{});
+
+    // Compute tile residues for predication.
+    auto m_max_coord = size<0>(shape_MNKL) - size<0>(gA) * get<0>(cta_coord);
+    auto n_max_coord = size<1>(shape_MNKL) - size<0>(gB) * get<1>(cta_coord);
+    auto k_residue = size<2>(shape_MNKL) - size<1>(gA) * size<2>(gA);
+    auto residue_mnk = make_tuple(m_max_coord, n_max_coord, k_residue);
+
+    // Allocate the tiled_mma and the accumulators for the (M,N) cta_tile.
+    TiledMma tiled_mma;
+    Tensor accum = partition_fragment_C(tiled_mma, take<0, 2>(cta_tile));
+    clear(accum);
+
+    auto k_tile_iter = make_coord_iterator(shape<2>(gA));
+    int k_tile_count = size<2>(gA);
+
+    // Perform the collective scoped MMA.
+    CollectiveMainloop collective_mma;
+    collective_mma(
+        accum,
+        gA,
+        gB,
+        accum,
+        k_tile_iter,
+        k_tile_count,
+        residue_mnk,
+        thread_idx,
+        smem_buf);
+
+    // Epilogue and write to out.
+    CollectiveEpilogue epilogue(params.epilogue);
+    epilogue(
+        shape_MNKL,
+        cta_tile,
+        cta_coord,
+        accum,
+        tiled_mma,
+        residue_mnk,
+        thread_idx,
+        smem_buf);
+  }
+};
+
+template <typename Element, bool KMajor>
+struct SimtCopyTraits {};
+
+template <typename Element>
+struct SimtCopyTraits<Element, true> {
+  using GmemTiledCopy = decltype(make_tiled_copy(
+      Copy_Atom<UniversalCopy<Element>, Element>{},
+      Layout<Shape<_32, _8>, Stride<_8, _1>>{},
+      Layout<Shape<_1, _1>>{}));
+  using SmemLayout = Layout<Shape<_128, _8>, Stride<_1, Int<128 + 1>>>;
+  using SmemCopyAtom = Copy_Atom<DefaultCopy, Element>;
+};
+
+template <typename Element>
+struct SimtCopyTraits<Element, false> {
+  using GmemTiledCopy = decltype(make_tiled_copy(
+      Copy_Atom<UniversalCopy<Element>, Element>{},
+      Layout<Shape<_32, _8>, Stride<_1, _32>>{},
+      Layout<Shape<_1, _1>>{}));
+  using SmemLayout = Layout<Shape<_128, _8>, Stride<_1, _128>>;
+  using SmemCopyAtom = Copy_Atom<DefaultCopy, Element>;
+};
+
+template <typename F>
+void dispatch_stride(bool k_major, int m, int k, F&& f) {
+  if (k_major) {
+    f(make_stride(k, Int<1>{}, m * k), std::true_type{});
+  } else {
+    f(make_stride(Int<1>{}, m, m * k), std::false_type{});
+  }
+}
+
+template <typename Element, typename F>
+void gather_mm(
+    int m,
+    int n,
+    int k,
+    int l,
+    bool a_transposed,
+    bool b_transposed,
+    const Element* A,
+    const Element* B,
+    const uint32_t* lhs_indices,
+    const uint32_t* rhs_indices,
+    Element* C,
+    F&& launch_kernel) {
+  auto problem_shape = make_shape(m, n, k, l);
+  auto dC = make_stride(m, Int<1>{}, m * n);
+  dispatch_stride(!a_transposed, m, k, [&](auto dA, auto k_major_a) {
+    dispatch_stride(b_transposed, n, k, [&](auto dB, auto k_major_b) {
+      using Accumulator =
+          std::conditional_t<(sizeof(Element) < 4), float, Element>;
+      using TileShape = Shape<_128, _128, _8>;
+      using DispatchPolicy = cutlass::gemm::MainloopSm70TwoStage;
+      using TiledMma = TiledMMA<
+          MMA_Atom<UniversalFMA<Accumulator, Element, Element, Accumulator>>,
+          Layout<Shape<_16, _16, _1>>>;
+
+      using CopyTraitsA = SimtCopyTraits<Element, k_major_a.value>;
+      using CopyTraitsB = SimtCopyTraits<Element, k_major_b.value>;
+
+      using CollectiveMainloop = cutlass::gemm::collective::CollectiveMma<
+          DispatchPolicy,
+          TileShape,
+          Element,
+          decltype(dA),
+          Element,
+          decltype(dB),
+          TiledMma,
+          typename CopyTraitsA::GmemTiledCopy,
+          typename CopyTraitsA::SmemLayout,
+          typename CopyTraitsA::SmemCopyAtom,
+          identity,
+          typename CopyTraitsB::GmemTiledCopy,
+          typename CopyTraitsB::SmemLayout,
+          typename CopyTraitsB::SmemCopyAtom,
+          identity>;
+
+      using CollectiveEpilogue = cutlass::epilogue::collective::DefaultEpilogue<
+          Element,
+          decltype(dC),
+          decltype(dC),
+          cutlass::epilogue::thread::
+              LinearCombination<Element, 1, Accumulator, Accumulator>,
+          cutlass::gemm::EpilogueDefault>;
+
+      using GemmKernel = GatherGemm<
+          decltype(problem_shape),
+          CollectiveMainloop,
+          CollectiveEpilogue>;
+      using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+
+      Gemm gemm;
+      typename Gemm::Arguments args{
+          problem_shape,
+          lhs_indices,
+          rhs_indices,
+          {A, dA, B, dB},
+          {{1.f, 0.f}, C, dC, C, dC}};
+
+      CHECK_CUTLASS_ERROR(gemm.initialize(args, nullptr));
+
+      auto* kernel = &cutlass::device_kernel<GemmKernel>;
+      void* kernel_params[] = {const_cast<Gemm::Params*>(&gemm.params())};
+      launch_kernel(
+          reinterpret_cast<void*>(kernel),
+          gemm.get_grid_shape(gemm.params()),
+          GemmKernel::get_block_shape(),
+          GemmKernel::SharedStorageSize,
+          kernel_params);
+    });
+  });
+}
+
+} // namespace cutlass_gemm
+
+namespace mlx::core {
+
+void cutlass_gather_mm(
+    bool a_transposed,
+    bool b_transposed,
+    const array& a,
+    const array& b,
+    const array& lhs_indices,
+    const array& rhs_indices,
+    array& out,
+    cu::CommandEncoder& encoder) {
+  int m = out.shape(-2);
+  int n = out.shape(-1);
+  int k = a.shape(-1);
+  int l = out.size() / (m * n);
+
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(lhs_indices);
+  encoder.set_input_array(rhs_indices);
+  encoder.set_output_array(out);
+
+  dispatch_float_types(out.dtype(), "gather_mm", [&](auto type_tag) {
+    using Element = cutlass_type_t<MLX_GET_TYPE(type_tag)>;
+    cutlass_gemm::gather_mm(
+        m,
+        n,
+        k,
+        l,
+        a_transposed,
+        b_transposed,
+        gpu_ptr<Element>(a),
+        gpu_ptr<Element>(b),
+        gpu_ptr<uint32_t>(lhs_indices),
+        gpu_ptr<uint32_t>(rhs_indices),
+        gpu_ptr<Element>(out),
+        [&](auto* kernel,
+            dim3 num_blocks,
+            dim3 block_dims,
+            uint32_t smem_bytes,
+            void** args) {
+          encoder.add_kernel_node_raw(
+              kernel, num_blocks, block_dims, {}, smem_bytes, args);
+        });
+  });
+}
+
+} // namespace mlx::core
@@ -0,0 +1,23 @@
+// Copyright © 2026 Apple Inc.
+
+#pragma once
+
+namespace mlx::core {
+
+namespace cu {
+class CommandEncoder;
+}
+
+class array;
+
+void cutlass_gather_mm(
+    bool a_transposed,
+    bool b_transposed,
+    const array& a,
+    const array& b,
+    const array& lhs_indices,
+    const array& rhs_indices,
+    array& out,
+    cu::CommandEncoder& encoder);
+
+} // namespace mlx::core
@@ -167,7 +167,7 @@ __global__ void gemv_gather(
 }

 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));
+  return (M == 1 && b_transposed) || (N == 1 && !a_transposed);
 }

 template <typename F>
@@ -236,7 +236,6 @@ void gemv(
            kernel,
            num_blocks_x,
            block_dims,
-            0,
            mat,
            vec,
            gpu_ptr<DataType>(out),
@@ -248,7 +247,6 @@ void gemv(
            kernel,
            dim3{num_blocks_x, batch_count},
            block_dims,
-            0,
            mat,
            vec,
            gpu_ptr<DataType>(out),
@@ -302,7 +300,6 @@ void gather_mv(
          kernel,
          dim3{num_blocks_x, batch_size},
          block_dims,
-          0,
          mat,
          vec,
          gpu_ptr<DataType>(out),
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Cheng	1091e3dd0a	Use uv in macOS CI	2026-05-06 15:41:43 +09:00
Cheng	80bcd1c658	[CUDA] Fix half type matmul in cutlass kernels (#3469 )	2026-05-06 08:35:53 +09:00
serenposh	1fdd4e23c2	Clearer error when shape dimension overflows int32 (#3425 ) Co-authored-by: Kanishk <kanishk.chores@gmail.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>	2026-05-05 09:53:36 +09:00
Pedro Cuenca	b43965925f	Define ST_F8_E8M0 (#3448 )	2026-05-05 09:22:23 +09:00
Abhilash Shankarampeta	0938db7e54	Add determinant and sign-log-determinant functions to mlx.core.linalg (#3416 ) Co-authored-by: Lucas Fernandes Martins <Lucas-Fernandes-Martins@users.noreply.github.com>	2026-05-05 09:06:23 +09:00
Irakli Salia	e8ebdebeeb	Add barrier to JACCL (#3459 )	2026-04-28 09:39:56 -07:00
Cheng	d7d0992d75	Reuse nightly build's ccache for release (#3458 )	2026-04-28 10:54:41 +09:00
Kimon N.	bdb6ff8881	Keep gguflib input-validation asserts active in release builds (#3436 )	2026-04-27 08:46:57 +09:00
Long Yixing	894c948773	[CUDA] Fix qmm_naive K-tail dispatch for FP quantized kernels (#3445 )	2026-04-27 08:40:14 +09:00
Angelos Katharopoulos	211e57be53	Bump minor (#3438 )	2026-04-22 11:09:30 -07:00
Cheng	c284e0a231	Enable swap for all CI building CUDA (#3437 )	2026-04-22 13:13:24 +09:00
Cheng	b9b1bfb9a5	Generate qmm implementaions with cmake (#3424 )	2026-04-22 13:11:55 +09:00
Cameron Churchwell	68cf2fddd8	Fix mx.prod vjp for complex types (#3433 )	2026-04-21 17:35:20 -07:00
Doğukan Veziroğlu	c594e6ec38	Fix use after move in Compiled primitive (#3427 )	2026-04-21 15:22:45 -07:00
Doğukan Veziroğlu	7d40a4fd5a	Throw meaningful error when Metal device is not found (#3428 )	2026-04-21 15:21:08 -07:00
Doğukan Veziroğlu	5f519ef6f9	Fix bytes_per_key truncation in random kernels (Metal + CUDA) (#3432 )	2026-04-21 15:15:11 -07:00
Angelos Katharopoulos	705c828feb	Fix synchronize for ThreadLocalStream (#3429 )	2026-04-20 11:29:49 -07:00
Cheng	b4ddf9b374	Fix flaky TestVmap.test_vmap_masked_scatter (#3421 )	2026-04-20 17:19:20 +09:00
Cheng	1f5a413a27	Make Scheduler::enqueue thread safe (#3423 )	2026-04-20 14:30:05 +09:00
Angelos Katharopoulos	a6222f53d5	Speed up NAX split-K by better tuning and routing and fix NAX addmm (#3422 ) I 'll merge now and comment with more benchmarks later since this also fixes two bugs so worst case we 'll do another tuning, it isn't like we won't need the functionality of this PR.	2026-04-19 18:05:39 -07:00
Cheng	fa4320d5fa	[CUDA] Handle residue k in qmm_naive (#3379 )	2026-04-18 13:30:07 +09:00
Long Yixing	859f22fbb0	[CUDA] GatherQMM matrix-matrix sm80/naive path (#3417 ) Co-authored-by: Cheng <git@zcbenz.com>	2026-04-18 10:59:47 +09:00
Cheng	d142de6a20	[CUDA] gather_mm (#3414 )	2026-04-17 16:53:44 +09:00
Angelos Katharopoulos	940ba473fe	Segmented mm nax kernel (#3419 )	2026-04-16 17:26:29 -07:00
Angelos Katharopoulos	8e649be4d0	Fix jaccl init bug (#3418 )	2026-04-16 01:23:35 -07:00
Cheng	dec6b4d10f	ThreadLocalStream in C++ (#3405 )	2026-04-15 15:46:11 -07:00
NeuralNoble	fd8e849e26	Document sort stability and NaN handling (#3400 )	2026-04-15 14:32:42 -07:00
Matias Insaurralde	50ae31241a	Validate safetensors data offsets against file boundaries (#3410 ) Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2026-04-15 14:30:55 -07:00
Dan Anderson	6cef1e995e	Validate safetensors data offsets (#3364 )	2026-04-15 00:52:42 -07:00
Cheng	57bcced8cb	Fixes for CUDA CI (#3413 )	2026-04-14 23:52:52 -07:00
Angelos Katharopoulos	4400504ad5	Jaccl refactor (#3412 )	2026-04-14 23:52:21 -07:00
jrp2014	1fa764fbec	Update nanobind version to v2.12.0 (#3396 )	2026-04-14 17:21:00 -07:00
Cheng	435f0b6cdb	Add clear_streams API for cleanup before exit (#3395 )	2026-04-14 18:41:32 +09:00
Cheng	520cea2bec	Avoid joining threads on exit (#3388 )	2026-04-11 09:22:34 +09:00
Clydingus	a33b791615	Fix int16 overflow in SDPA NAX mask indexing for KV sequences > 32K (#3361 ) Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2026-04-10 00:01:47 -07:00
Cameron Churchwell	d6d9b24801	Conjugate VJP and JVP support (#3386 )	2026-04-09 15:04:46 -07:00
Daniil Seredkin	8332e228e4	Fix test "test get streams" missing initialization (#3376 )	2026-04-09 08:29:04 +09:00
Cheng	4403165843	[CUDA] Thread safety (#3367 )	2026-04-09 08:18:00 +09:00
Shantanu Suryawanshi	a8776b7bbd	Fix: Correct cross-attention query routing in Post-LN TransformerDecoderLayer (#3382 )	2026-04-07 09:16:12 -07:00
Doğukan Veziroğlu	b98831ad0e	fix: fail build when Metal compiler header resolution fails (#3332 )	2026-04-06 12:49:25 -07:00
Long Yixing	d025111b1d	[CUDA] Add GatherQMM for quantized gather matmul (#3321 )	2026-04-06 12:48:18 -07:00
Harrison Powers	9239808225	Fix CMake finding wrong Python during pip install (#3375 )	2026-04-06 12:32:16 -07:00
Angelos Katharopoulos	6a9a121d09	Add a convenience for making local streams in python (#3355 )	2026-04-02 18:43:02 -07:00
Christophe Prat	befe42d303	Add printoptions (#3333 )	2026-04-01 22:24:48 -07:00
Valentin Roussellet	80a1c206f9	Use `metal` as the front-end for the metal linker (#3354 )	2026-04-01 16:52:07 -07:00
Angelos Katharopoulos	b0748ad8de	Fix regression in array creation (#3353 )	2026-04-01 11:30:36 -07:00
Cheng	2ffafe07f4	[CUDA] 3/5/6-bit quants for qmm_naive (#3352 )	2026-04-01 20:13:01 +09:00
Cheng	5e2c44259f	Make CommandEncoder thread local (#3348 )	2026-04-01 18:42:49 +09:00
Cheng	1c9ee2f655	[CUDA] Fallback QMM (#3315 )	2026-04-01 12:41:26 +09:00
Long Yixing	7cd73c4202	[Metal] Support sorting complex numbers (#3314 )	2026-04-01 12:40:50 +09:00
declanhealy2	2105df91da	Add fftfreq, rfftfreq and scalar axes for fftshift/ifftshift (#3298 )	2026-03-31 18:29:16 -07:00
Angelos Katharopoulos	1944cf67a2	Add vmap for BroadcastAxes (#3344 )	2026-03-31 17:08:56 -07:00
Cheng	939e425c7a	Decouple CommandEncoder from Device (#3316 )	2026-04-01 08:51:17 +09:00
Angelos Katharopoulos	8439b1f501	Fix use after move (#3343 )	2026-03-31 10:37:40 -07:00
dependabot[bot]	117b4f1806	Bump actions/deploy-pages from 4 to 5 (#3334 ) Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>	2026-03-31 09:04:19 +09:00
Cheng	66f58032dc	Remove no longer needed const_cast (#3325 )	2026-03-31 08:10:49 +09:00
Kellen Sun	8a6d28713c	Fix np bfloat16 misinterpreted as complex (#3146 ) Co-authored-by: Cheng <git@zcbenz.com>	2026-03-31 08:04:55 +09:00
Long Yixing	0ff1115a46	[CUDA] Implement BlockMaskedMM (#3299 ) Co-authored-by: Cheng <git@zcbenz.com>	2026-03-27 06:57:26 +09:00
Cheng	df7f7db943	Make each thread have its own default stream (#3281 )	2026-03-25 15:48:49 +09:00
Sheldon Aristide	57c813f042	Add norm parameter to FFT transforms (#3287 ) Co-authored-by: Cheng <git@zcbenz.com>	2026-03-25 13:27:40 +09:00
Long Yixing	f8eda2c61b	[CUDA] support sorting complex numbers (#3286 )	2026-03-25 12:35:02 +09:00
Cheng	282174dd03	Manage Metal objects with smart pointers (#3282 )	2026-03-25 11:19:20 +09:00
Pranav Hari	bd200d6267	Add output_shapes for AddMM (#3262 ) Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2026-03-25 10:46:15 +09:00
Cheng	d01b83dfe7	Use nb::ndarray for checking arrays (#3283 )	2026-03-25 10:44:54 +09:00
Sheldon Aristide	1b1c56352a	Fix moved-from shape bug in broadcast_arrays causing vmap bus error (#3310 )	2026-03-24 17:02:31 -07:00
Robert Johansson	e18d4e97f6	Fix vmap + floor_divide: preserve integer dtype (#3292 ) Co-authored-by: Robert Johansson <robert@Mac-Mini-KI.lan> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Angelos Katharopoulos <katharas@gmail.com> Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2026-03-25 08:10:37 +09:00
Ronan Collobert	9ab3913567	logo files (#3308 )	2026-03-24 15:08:06 -07:00
Sheldon Aristide	81530c261b	Implement Pad::vmap (#3304 ) Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2026-03-24 15:02:31 -07:00
LongYinan	604c825538	Fix stale transform copy-chain leaks (#3290 )	2026-03-24 14:15:23 -07:00
LongYinan	e40ada3fe2	[Metal] Fix depthwise conv 1D kernel name for large variant (#3289 )	2026-03-23 16:20:13 -07:00
Ziqiao-git	38ad257088	[Metal][Performance]: Add split-K for quantized matmul (small M) (#3120 )	2026-03-20 20:15:48 -07:00
Cheng	70a0da6fca	Use thread local storage for frontend compile cache (#3280 )	2026-03-20 07:44:45 +09:00
Long Yixing	82809ebd12	Fix sort NaN handling for float16 and bfloat16 (#3269 )	2026-03-19 15:19:41 -07:00
AN Long	5fa1a8d59f	Support indexing with any type which implmented `__index__` (#3210 )	2026-03-19 15:19:08 -07:00
Cheng	21c11fc9b0	Create default random key lazily (#3278 )	2026-03-19 20:22:52 +09:00
Cheng	e1cbac9cf4	[CUDA] Search system-installed CUDA toolkit for headers (#3277 )	2026-03-19 20:09:05 +09:00
Cheng	c8292ea11c	Merge DeviceStream into CommandEncoder (#3264 )	2026-03-19 19:39:30 +09:00
Angelos Katharopoulos	45af0df90b	Fix repr of conv layers (#3275 )	2026-03-18 22:47:38 -07:00
Cheng	dbfbc0f65a	[CUDA] fp and int4 quants for qmm_sm80 (#3268 )	2026-03-19 09:38:55 +09:00
Cheng	75f74ea9bc	Fix building with CUDA toolkit 13.2 (#3273 )	2026-03-19 08:31:44 +09:00
Jagrit Digani	b41b349b67	Nax Refactor (#3271 )	2026-03-18 10:26:49 -07:00
Angelos Katharopoulos	7bc61cceed	Slice update with operation (#3266 )	2026-03-18 06:18:02 -07:00
Ihar Hrachyshka	e353be8235	tests: harden memory leak check in test_siblings_without_eval (#3088 ) Signed-off-by: Ihar Hrachyshka <ihar.hrachyshka@gmail.com>	2026-03-17 16:03:10 +09:00
Cheng	1e855446b2	[CUDA] Pipelined QMM (#3255 )	2026-03-17 07:10:12 +09:00
mm65x	f226eeec9e	Fix nn.GRU skipping bhn bias when hidden is None (#3252 ) Co-authored-by: mm65x <mm65x@users.noreply.github.com>	2026-03-16 13:28:14 -07:00
mm65x	505fc9850d	Fix comparison op JVP returning bool tangents instead of input dtype (#3253 )	2026-03-16 10:57:28 -07:00
Thomas Schranz	ea91bd02cf	update requirements for Macbook Neo (#3257 )	2026-03-16 04:33:09 -07:00
Lik Xun Yuan (Lx)	1d44d913e6	docs: fix PyTorch to MLX conversion example (#3265 )	2026-03-16 04:20:12 -07:00
Long Yixing	0bdbfdb838	[CUDA] Implement MaskedScatter (#3151 )	2026-03-15 10:33:55 +09:00
Lucas Newman	5d1700493a	[CUDA] Add FFT support (#3243 )	2026-03-14 21:02:19 +09:00
Valentin Roussellet	b0564a9112	Fix crashes in multi-threaded process teardown (#3167 )	2026-03-12 21:45:06 -07:00
Daniel Hiltgen	7adfc83c7d	win: re-enable and fix cuDNN performance (#3242 )	2026-03-13 09:41:59 +09:00
Angelos Katharopoulos	0358c602c7	Bump (#3244 )	2026-03-11 23:57:22 -07:00
Cheng	ce45c52505	[CUDA] Use qmv kernel for fp quantizations (#3239 )	2026-03-12 07:25:17 +09:00
Jagrit Digani	0879a6acba	Add initial tuning for M5 pro and max (#3211 )	2026-03-11 14:05:43 -07:00
Long Yixing	a9573f92f6	[CUDA] Implement SegmentedMM (#3238 )	2026-03-11 13:31:43 -07:00
Cheng	1c2d7041ab	Remove quantized_utils.cuh (#3237 )	2026-03-11 19:45:51 +09:00
Daniel Hiltgen	fd6d304b3a	win: fix cuda build (#3204 )	2026-03-11 12:58:04 +09:00
Anastasiia Filippova	e1e1399e1b	Hybrid sharding (#3194 )	2026-03-10 11:47:25 +01:00
Cheng	9d03a1b0d9	[CUDA] Support 3/5/6-bit quants in QMV (#3236 )	2026-03-10 19:09:48 +09:00
Cheng	8d022bcb86	Remove custom fp4/fp8 classes (#3212 )	2026-03-10 16:08:01 +09:00
Michelle DiMarco	d2702a4fc1	Fix non-strict module update with extra weights (#3214 ) Co-authored-by: Michelle DiMarco <m_dimarco@apple.com>	2026-03-09 22:03:50 -07:00
Dan Anderson	6ac5280db4	Fix assigning bool to float16/bfloat16 (#3229 ) Co-authored-by: KD2YCU <me@kd2ycu.com> Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2026-03-09 21:50:05 -07:00
Dan Anderson	572e0a4ac3	Validate dims in rope (#3230 ) Co-authored-by: KD2YCU <me@kd2ycu.com>	2026-03-09 21:48:45 -07:00
Dan Anderson	9bbd375eec	Fix return value in einsum_path for simple contractions (#3232 ) Co-authored-by: KD2YCU <me@kd2ycu.com> Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2026-03-09 21:48:26 -07:00
Dan Anderson	a25399cbd4	Validate num_splits in split (#3234 ) Co-authored-by: KD2YCU <me@kd2ycu.com>	2026-03-09 21:47:08 -07:00
Cheng	5a347b2ec8	[CUDA] Faster compilation and batch support in QMV (#3213 )	2026-03-10 13:45:10 +09:00
Dan Anderson	db487f3649	PR #3220 LayerNorm VJP returns zeros_like(weight) instead of zeros_like(bias placeholder) (#3231 ) Co-authored-by: KD2YCU <me@kd2ycu.com>	2026-03-09 17:06:12 -07:00
Dan Anderson	8f5ff2ea41	PR#3226 Fix (#3227 ) Co-authored-by: KD2YCU <me@kd2ycu.com>	2026-03-09 15:06:33 -07:00
Angelos Katharopoulos	d06c3c8936	Improve mlx.distributed_config (#3199 )	2026-03-09 13:17:51 -07:00
Long Yixing	be872ebdef	[CUDA] implement Hadamard transform (#3179 )	2026-03-05 09:34:19 +01:00
Cheng	3b3590bf5f	[CUDA] Use fp16 accumulation for 4-bit quant in GEMV (#3197 )	2026-03-05 07:58:23 +09:00
Cheng	3c565437a5	[CUDA] Quantized GEMV (#3180 )	2026-03-04 08:59:31 +09:00
dependabot[bot]	9eef9f1774	Bump actions/upload-artifact from 6 to 7 (#3188 ) Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>	2026-03-04 08:26:24 +09:00
dependabot[bot]	e320a2adc2	Bump actions/download-artifact from 7 to 8 (#3189 ) Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>	2026-03-04 08:24:12 +09:00
willem adnet	8cd377b7db	Add the bartlett function (#3155 )	2026-03-03 11:40:54 -08:00
Christophe Prat	f145ece976	Fix/missing libs in docs (#3190 )	2026-03-03 09:53:18 -08:00
AN Long	1ce7118303	Fix ref leak in mx.save/load with file like object (#3187 )	2026-03-02 16:55:29 -08:00
Anastasiia Filippova	72e04f7fb7	[CUDA] Fsdp (easy) (#3130 )	2026-03-01 23:29:09 +01:00
Cheng	6482d13dd3	Skip Hopper-only kernels in CI (#3184 )	2026-02-27 23:22:08 -08:00
Angelos Katharopoulos	1e3736b19d	Bump the patch version (#3185 )	2026-02-27 23:21:37 -08:00
Angelos Katharopoulos	365d6f29b4	Bump the minor version (#3183 )	2026-02-27 14:02:03 -08:00
Angelos Katharopoulos	d7a553c536	Enable passing in a GPU architecture string via env var (#3176 )	2026-02-27 11:37:53 -08:00
Robert Johansson	c8536f5248	Fix compile_fuse broadcast split aliasing bug (#3166 ) Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2026-02-26 18:17:28 -08:00
Cheng	0c8107ce8a	[CUDA] Heuristics for Hopper QMM (#3173 )	2026-02-27 09:22:10 +09:00
Angelos Katharopoulos	5c4abd2f06	JACCL refactor and small update (#3174 )	2026-02-26 13:56:19 -08:00
Anastasiia Filippova	4e00919e5c	[CUDA][NCCL] group split (#3172 )	2026-02-26 09:26:20 +01:00
Cheng	6ec0192270	Enable setting thread block cluster for Hopper and later (#3168 )	2026-02-26 16:15:16 +09:00
willem adnet	a8ba5ac3e0	Implement mlx.core.blackman (#3136 )	2026-02-25 13:42:40 -08:00
Cheng	6304c285d3	[CUDA] FPxINT quantized matmul for Hopper (#3160 )	2026-02-25 09:10:18 +09:00
Awni Hannun	cb198268d5	[Metal] Fix event leak (#3159 )	2026-02-23 19:50:48 -08:00
Gleb Sterkin	1d8d693d08	[Metal] Add implicit matmul pathway for mx.conv3d (#3147 ) Co-authored-by: Gleb Sterkin <g_sterkin@apple.com> Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2026-02-23 17:52:50 -08:00
Kellen Sun	d4c81062ad	[Metal] Fix 32-bit integer overflow in conv3d unfold kernel (#3143 )	2026-02-19 10:01:22 -08:00
Alex Skryl	f2f2d16451	Export: preserve Dtype state values in export callback arguments (#3145 ) Co-authored-by: Awni Hannun <awni@apple.com>	2026-02-19 08:07:28 -08:00
Awni Hannun	daf18e76ca	Fix fence synchronization accross command buffers (#3144 )	2026-02-18 20:21:26 -08:00
Anastasiia Filippova	06305022ab	Tensor scale nvfp4 (#3022 )	2026-02-18 11:19:26 +01:00
willem adnet	360639c2df	Add the hamming window function (#3135 )	2026-02-17 00:56:05 -08:00
willem adnet	3bbe87e6dc	Add hanning window function (#3124 )	2026-02-16 09:44:49 -08:00
vskiwi	e226af720e	Propagate quantization mode in quantized layers (#3133 )	2026-02-15 18:33:13 -08:00
Awni Hannun	43f4a74826	Manage stream placement in import function (#3127 )	2026-02-15 06:17:06 -08:00
Angelos Katharopoulos	c184262d29	Fix donation in sdpa vector (#3121 )	2026-02-12 10:46:21 -08:00
Cheng	72e94c81e1	[CUDA] Attention sinks in cuDNN SDPA (#3118 )	2026-02-11 16:46:39 +09:00
Awni Hannun	4c86c1e55a	Fix precision in Metal fused attention (#3119 )	2026-02-10 14:18:29 -08:00
Anastasiia Filippova	be52cf660b	register pressure (#3116 )	2026-02-10 14:17:28 -08:00
Cheng	54bb3eea42	[CUDA] Use cuDNN SDPA for decoding when using fixed-size KV cache (#3113 )	2026-02-10 09:15:45 +09:00
Anastasiia Filippova	5e018de4e5	Quantize module to QQLinear (#3106 )	2026-02-09 14:35:17 -08:00
Cheng	9cd4b9be91	[CUDA] Set current device before allocating memory (#3110 )	2026-02-08 19:04:57 +09:00
Cheng	566bc16b7c	Cleanup test_fast_sdpa.py (#3112 )	2026-02-08 19:04:24 +09:00
Awni Hannun	8fe1d09207	Fix residency set with user provided buffer (#3108 )	2026-02-06 16:38:36 -08:00
Ronan Collobert	ef3fbc60a3	is_available() should check the device index too (#3107 )	2026-02-06 13:02:04 -08:00
Angelos Katharopoulos	69fd3fa9b1	Patch bump (#3102 )	2026-02-06 09:15:22 -08:00
Awni Hannun	185b06d9ef	Patch for multi device CUDA (#3100 )	2026-02-05 17:33:51 -08:00
Manuel Candales	90e38f7b93	Fix qmv_impl for small N (#3096 )	2026-02-05 17:33:36 -08:00
Angelos Katharopoulos	ceea571490	JACCL update (#3094 )	2026-02-05 15:16:07 -08:00
Awni Hannun	99ca62c4d3	Fix 2pass sdpa on < M2 (#3099 )	2026-02-05 08:51:29 -08:00
Awni Hannun	206cf07e5b	Fix non simd f16 build (#3097 )	2026-02-05 07:04:02 -08:00
Jesse Gross	f47729c0d8	Disable managed memory on WSL when concurrentManagedAccess is not supported (#3095 )	2026-02-05 10:58:49 +09:00
Awni Hannun	b9b672250e	patch (#3093 )	2026-02-03 07:24:30 -08:00
Awni Hannun	adcbb91a9e	Fix for NAX overflow. (#3092 )	2026-02-02 18:54:01 -08:00
Awni Hannun	b56782be52	[Metal] Tune splitk gemm dispatch conditions and partition sizes (#3087 )	2026-02-02 08:45:09 -08:00
Cheng	8ef539522c	Fix failing python tests on Windows (#3076 )	2026-01-30 17:50:18 +09:00
Cheng	212077f163	Fallback to pinned host memory when managed memory is not supported (#3075 )	2026-01-30 13:18:41 +09:00
Awni Hannun	cc6e4eebad	Fix nax condition for iphone (#3083 )	2026-01-29 13:30:37 -08:00
Awni Hannun	fcbdd05022	More useful error for large indices (#3079 )	2026-01-29 13:02:39 -08:00
atharva	590b4f1c16	Fix ALiBi slopes for non-power-of-2 num_heads (#3071 )	2026-01-29 07:23:11 -08:00
Anri Lombard	0c6a895ed7	Use lower-right causal mask alignment consistently (#2967 ) Co-authored-by: Awni Hannun <awni.hannun@gmail.com>	2026-01-28 17:15:14 -08:00
stef	c86a9bced1	[Docs] Simple example of using MLX distributed (#2973 ) Co-authored-by: Awni Hannun <awni@apple.com>	2026-01-28 17:14:56 -08:00
Awni Hannun	12e386f308	Tune CUDA gaph sizes on B200 and H100 (#3077 )	2026-01-28 17:14:44 -08:00
Cheng	2ac18eddb9	[CUDA] Fallback Event impl when there is no hardware cpu/gpu coherency (#3070 )	2026-01-28 10:43:22 +09:00
Awni Hannun	b537b3685f	patch (#3074 )	2026-01-28 10:42:55 +09:00