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
339 changed files with 20103 additions and 6469 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(
@@ -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}"
+            )
@@ -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
@@ -32,7 +32,7 @@ are the CPU and GPU.
   install

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

   usage/quick_start
@@ -78,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::

@@ -0,0 +1,12 @@
+Print Options
+===============
+
+.. currentmodule:: mlx.core
+
+.. autosummary::
+  :toctree: _autosummary
+
+  PrintOptions
+  set_printoptions
+  printoptions
+  get_printoptions
@@ -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
@@ -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;
  }
@@ -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) {
@@ -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);
@@ -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))
@@ -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)
@@ -168,9 +168,6 @@ CudaAllocator::CudaAllocator()
  free_limit_ = total_memory_ - memory_limit_;
  max_pool_size_ = memory_limit_;

-  int curr;
-  CHECK_CUDA_ERROR(cudaGetDevice(&curr));
-
  int device_count = gpu::device_count();
  free_streams_.resize(device_count);
  mem_pools_.resize(device_count);
@@ -181,7 +178,6 @@ CudaAllocator::CudaAllocator()
      CHECK_CUDA_ERROR(cudaDeviceGetDefaultMemPool(&mem_pools_[i], i));
    }
  }
-  CHECK_CUDA_ERROR(cudaSetDevice(curr));
 }

 Buffer
@@ -223,6 +219,7 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
      if (device == -1) {
        data = unified_malloc(size);
      } else {
+        cu::device(device).make_current();
        if (mem_pools_[device]) { // supports memory pools
          CHECK_CUDA_ERROR(cudaMallocAsync(&data, size, stream));
        } else {
@@ -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)>
@@ -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) {
@@ -50,49 +54,20 @@ Device::Device(int device) : device_(device) {
      &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()) {
@@ -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_;
  }
@@ -183,13 +201,12 @@ class Device {
  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
@@ -213,7 +213,7 @@ auto check_gpu_coherency() {
  return coherency;
 }

-AtomicEvent::AtomicEvent() {
+AtomicEvent::AtomicEvent(Device& d) {
  void* buf;
  cudaError_t (*cuda_free)(void*);
  // There are 3 kinds of systems we are implementing for:
@@ -223,6 +223,7 @@ AtomicEvent::AtomicEvent() {
  //    => 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)));
@@ -347,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,7 +54,7 @@ class CudaEvent {
 // CudaEvent so the latter should always be preferred when possible.
 class AtomicEvent {
 public:
-  AtomicEvent();
+  AtomicEvent(Device& d);

  void wait(uint32_t value);
  void wait(cudaStream_t stream, uint32_t value);
@@ -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&) {
@@ -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),
@@ -22,4 +22,18 @@ void cutlass_grouped_gemm_unaligned(
    array& out,
    cu::CommandEncoder& encoder);

+void cutlass_segmented_mm(
+    bool a_transposed,
+    int lda,
+    bool b_transposed,
+    int ldb,
+    int num_segments,
+    int M,
+    int N,
+    const array& a,
+    const array& b,
+    const array& segments,
+    array& out,
+    cu::CommandEncoder& encoder);
+
 } // namespace mlx::core
@@ -93,6 +93,50 @@ __global__ void prepare_grouped_mm_data(
  }
 }

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

 namespace {
@@ -158,10 +202,11 @@ class GemmGroupedEncoder
    : public cutlass::gemm::device::GemmGrouped<GemmKernel> {
 public:
  void encode(cu::CommandEncoder& encoder) {
-    encoder.add_kernel_node(
+    encoder.add_kernel_node_ex(
        cutlass::Kernel<GemmKernel>,
        {static_cast<uint32_t>(this->params_.threadblock_count), 1, 1},
        {GemmKernel::kThreadCount, 1, 1},
+        {},
        sizeof(typename GemmKernel::SharedStorage),
        this->params_);
  }
@@ -200,7 +245,7 @@ void grouped_gemm_v2(
          LayoutB,
          cutlass::ComplexTransform::kNone,
          GemmConfiguration::kAlignmentAB,
-          typename GemmConfiguration::Element,
+          typename GemmConfiguration::Accumulator,
          cutlass::layout::RowMajor,
          typename GemmConfiguration::Accumulator,
          typename GemmConfiguration::OpClass,
@@ -280,6 +325,11 @@ void cutlass_grouped_gemm_unaligned(
    const array& indices,
    array& out,
    cu::CommandEncoder& encoder) {
+  if (group_count > 1024) {
+    throw std::runtime_error(
+        "[gather_mm] Group count can not be larger than 1024.");
+  }
+
  int K = a.shape(-1);
  int N = b.shape(-1);

@@ -309,10 +359,11 @@ void cutlass_grouped_gemm_unaligned(

  encoder.set_input_array(indices);
  encoder.set_output_array(gemm_args);
-  encoder.add_kernel_node(
+  encoder.add_kernel_node_ex(
      cu::prepare_grouped_mm_data<N_READS>,
      num_blocks,
      block_dims,
+      {},
      group_count * sizeof(uint32_t), // sizeof(cum_histo)
      gpu_ptr<uint32_t>(indices),
      indices.size(),
@@ -355,4 +406,85 @@ void cutlass_grouped_gemm_unaligned(
      encoder);
 }

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

 #include "mlx/backend/common/compiled.h"
+#include "mlx/backend/common/slicing.h"
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/jit_module.h"
 #include "mlx/backend/cuda/kernel_utils.cuh"
 #include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/gpu/scan.h"
+#include "mlx/backend/gpu/slicing.h"
 #include "mlx/dtype_utils.h"
 #include "mlx/primitives.h"

@@ -23,6 +26,8 @@ namespace mlx::core {
 namespace {

 constexpr const char* g_scatter_ops[] = {"Max", "Min", "Sum", "Prod", "Assign"};
+constexpr const char* g_slice_ops[] =
+    {"Maximum", "Minimum", "Add", "Multiply", ""};

 void append_indices_arg(
    cu::KernelArgs& args,
@@ -130,7 +135,8 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {

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

 void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
@@ -232,7 +238,8 @@ void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
  encoder.set_output_array(out);
  auto kernel = mod.get_kernel(kernel_name);
  auto [num_blocks, block_dims] = get_launch_args(upd, large);
-  encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
+  encoder.add_kernel_node_raw(
+      kernel, num_blocks, block_dims, {}, 0, args.args());
 }

 void GatherAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
@@ -322,7 +329,8 @@ void GatherAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
  encoder.set_output_array(out);
  auto kernel = mod.get_kernel(kernel_name);
  auto [num_blocks, block_dims] = get_launch_args(idx, large);
-  encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
+  encoder.add_kernel_node_raw(
+      kernel, num_blocks, block_dims, {}, 0, args.args());
 }

 void ScatterAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
@@ -428,7 +436,250 @@ void ScatterAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
  encoder.set_output_array(out);
  auto kernel = mod.get_kernel(kernel_name);
  auto [num_blocks, block_dims] = get_launch_args(idx, large);
-  encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args());
+  encoder.add_kernel_node_raw(
+      kernel, num_blocks, block_dims, {}, 0, args.args());
+}
+
+void MaskedScatter::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("MaskedScatter::eval_gpu");
+  assert(inputs.size() == 3);
+
+  const array& dst = inputs[0];
+  const array& mask = inputs[1];
+  const array& src = inputs[2];
+
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  const size_t total = mask.size();
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+  if (total == 0) {
+    return;
+  }
+
+  array mask_flat = flatten_in_eval(mask, 1, -1, s);
+  if (gpu_ptr<void>(mask_flat) != gpu_ptr<void>(mask)) {
+    encoder.add_temporary(mask_flat);
+  }
+  if (!mask_flat.flags().row_contiguous) {
+    mask_flat = contiguous_copy_gpu(mask_flat, s);
+    encoder.add_temporary(mask_flat);
+  }
+
+  array scatter_offsets(mask_flat.shape(), int32, nullptr, {});
+  scatter_offsets.set_data(cu::malloc_async(scatter_offsets.nbytes(), encoder));
+  encoder.add_temporary(scatter_offsets);
+
+  scan_gpu_inplace(
+      mask_flat,
+      scatter_offsets,
+      Scan::Sum,
+      /* axis= */ 1,
+      /* reverse= */ false,
+      /* inclusive= */ false,
+      s);
+
+  const size_t batch_count = mask.shape(0);
+  const size_t mask_batch_size = mask_flat.size() / batch_count;
+  const size_t src_batch_size = src.size() / src.shape(0);
+  bool large = total > INT32_MAX || src.size() > INT32_MAX;
+  bool vectorized = src.flags().row_contiguous && dst.flags().row_contiguous;
+  constexpr int kMaskedScatterVecSize = 16;
+  constexpr int kMaskedScatterVecBlockDim = 256;
+
+  std::string module_name =
+      fmt::format("masked_scatter_{}", dtype_to_string(out.dtype()));
+  cu::JitModule& mod = cu::get_jit_module(s.device, module_name, [&]() {
+    std::vector<std::string> kernel_names;
+    for (int src_contiguous = 0; src_contiguous <= 1; ++src_contiguous) {
+      for (int dst_contiguous = 0; dst_contiguous <= 1; ++dst_contiguous) {
+        for (int use_large = 0; use_large <= 1; ++use_large) {
+          kernel_names.push_back(
+              fmt::format(
+                  "mlx::core::cu::masked_scatter<{}, {}, {}, {}>",
+                  dtype_to_cuda_type(out.dtype()),
+                  src_contiguous ? "true" : "false",
+                  dst_contiguous ? "true" : "false",
+                  use_large ? "int64_t" : "int32_t"));
+        }
+      }
+    }
+    for (int use_large = 0; use_large <= 1; ++use_large) {
+      kernel_names.push_back(
+          fmt::format(
+              "mlx::core::cu::masked_scatter_vec_contiguous<{}, {}, {}>",
+              dtype_to_cuda_type(out.dtype()),
+              use_large ? "int64_t" : "int32_t",
+              kMaskedScatterVecSize));
+    }
+    return std::make_tuple(false, jit_source_scatter, std::move(kernel_names));
+  });
+
+  cu::KernelArgs args;
+  args.append(dst);
+  args.append(mask_flat);
+  args.append(scatter_offsets);
+  args.append(src);
+  args.append(out);
+  if (large) {
+    args.append<int64_t>(mask_flat.size());
+    args.append<int64_t>(src_batch_size);
+    args.append<int64_t>(mask_batch_size);
+  } else {
+    args.append<int32_t>(mask_flat.size());
+    args.append<int32_t>(src_batch_size);
+    args.append<int32_t>(mask_batch_size);
+  }
+  if (!vectorized) {
+    args.append_ndim(dst.shape());
+    args.append_ndim(dst.strides());
+    args.append<int32_t>(dst.ndim());
+    args.append_ndim(src.shape());
+    args.append_ndim(src.strides());
+    args.append<int32_t>(src.ndim());
+  }
+
+  encoder.set_input_array(dst);
+  encoder.set_input_array(mask_flat);
+  encoder.set_input_array(scatter_offsets);
+  encoder.set_input_array(src);
+  encoder.set_output_array(out);
+
+  std::string kernel_name = vectorized
+      ? fmt::format(
+            "mlx::core::cu::masked_scatter_vec_contiguous<{}, {}, {}>",
+            dtype_to_cuda_type(out.dtype()),
+            large ? "int64_t" : "int32_t",
+            kMaskedScatterVecSize)
+      : fmt::format(
+            "mlx::core::cu::masked_scatter<{}, {}, {}, {}>",
+            dtype_to_cuda_type(out.dtype()),
+            src.flags().row_contiguous ? "true" : "false",
+            dst.flags().row_contiguous ? "true" : "false",
+            large ? "int64_t" : "int32_t");
+  auto kernel = mod.get_kernel(kernel_name);
+  auto [num_blocks, block_dims] = vectorized
+      ? get_launch_args(
+            mask_flat, large, kMaskedScatterVecSize, kMaskedScatterVecBlockDim)
+      : get_launch_args(mask_flat, large);
+  encoder.add_kernel_node_raw(
+      kernel, num_blocks, block_dims, {}, 0, args.args());
+}
+
+void SliceUpdate::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("SliceUpdate::eval_gpu");
+  assert(inputs.size() == 2);
+  if (out.size() == 0) {
+    return;
+  }
+
+  auto& in = inputs[0];
+  auto& upd = inputs[1];
+
+  if (upd.size() == 0) {
+    out.copy_shared_buffer(in);
+    return;
+  }
+
+  auto ctype = in.flags().contiguous && in.size() == in.data_size()
+      ? CopyType::Vector
+      : CopyType::General;
+  copy_gpu(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 for None reduce type
+  if (reduce_type_ == SliceUpdate::None) {
+    copy_gpu_inplace(
+        /* const array& src = */ upd,
+        /* array& dst = */ out,
+        /* const Shape& data_shape = */ upd.shape(),
+        /* const Strides& i_strides = */ upd.strides(),
+        /* const Strides& o_strides = */ out_strides,
+        /* int64_t i_offset = */ 0,
+        /* int64_t o_offset = */ data_offset,
+        /* CopyType ctype = */ CopyType::GeneralGeneral,
+        /* const Stream& s = */ stream());
+    return;
+  }
+
+  auto [shape, strides] =
+      collapse_contiguous_dims(upd.shape(), {upd.strides(), out_strides});
+  int nwork = 1;
+  if (shape.back() % 4 == 0) {
+    nwork = 4;
+  } else if (shape.back() % 2 == 0) {
+    nwork = 2;
+  }
+
+  const char* op_name = g_slice_ops[reduce_type_];
+  auto [ds, rc, cc] = check_contiguity(shape, strides[1]);
+  bool upd_contiguous = upd.flags().row_contiguous;
+  bool upd_scalar = upd.data_size() == 1;
+  bool out_contiguous = rc;
+  bool large = upd.size() > INT32_MAX;
+  std::string module_name =
+      fmt::format("slice_update_{}_{}", op_name, dtype_to_string(out.dtype()));
+
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  cu::JitModule& mod = cu::get_jit_module(s.device, module_name, [&]() {
+    std::vector<std::string> kernel_names;
+    for (int out_c = 0; out_c <= 1; ++out_c) {
+      for (int upd_c = 0; upd_c <= 1; ++upd_c) {
+        for (int upd_s = 0; upd_s <= 1; ++upd_s) {
+          for (int large = 0; large <= 1; ++large) {
+            for (int nwork = 1; nwork <= 16; nwork *= 2) {
+              kernel_names.push_back(
+                  fmt::format(
+                      "mlx::core::cu::slice_update_op<{}, {}, mlx::core::cu::{}, {}, {}, {}, {}>",
+                      dtype_to_cuda_type(out.dtype()),
+                      large ? "int64_t" : "int32_t",
+                      op_name,
+                      out_c ? "true" : "false",
+                      upd_c ? "true" : "false",
+                      upd_s ? "true" : "false",
+                      nwork));
+            }
+          }
+        }
+      }
+    }
+    return std::make_tuple(
+        false, jit_source_slice_update, std::move(kernel_names));
+  });
+
+  cu::KernelArgs args;
+  args.append(upd);
+  args.append(out);
+  args.append<int64_t>(upd.size());
+  args.append_ndim(shape);
+  args.append_ndim(strides[0]);
+  args.append<int32_t>(shape.size());
+  args.append_ndim(strides[1]);
+  args.append<int64_t>(data_offset);
+
+  encoder.set_input_array(upd);
+  encoder.set_output_array(out);
+
+  std::string kernel_name;
+  kernel_name = fmt::format(
+      "mlx::core::cu::slice_update_op<{}, {}, mlx::core::cu::{}, {}, {}, {}, {}>",
+      dtype_to_cuda_type(out.dtype()),
+      large ? "int64_t" : "int32_t",
+      op_name,
+      out_contiguous,
+      upd_contiguous,
+      upd_scalar,
+      nwork);
+
+  auto kernel = mod.get_kernel(kernel_name);
+  auto [num_blocks, block_dims] = get_launch_args(upd, large, nwork);
+  encoder.add_kernel_node_raw(
+      kernel, num_blocks, block_dims, {}, 0, args.args());
 }

 } // namespace mlx::core
@@ -9,6 +9,7 @@
 #include <cstdlib>
 #include <filesystem>
 #include <fstream>
+#include <shared_mutex>

 #include <fmt/format.h>
 #include <nvrtc.h>
@@ -26,6 +27,25 @@ void check_nvrtc_error(const char* name, nvrtcResult err) {
  }
 }

+// Return the default path to CUDA toolkit.
+const std::filesystem::path& default_cuda_toolkit_path() {
+#if defined(_WIN32)
+  static auto cached_path = []() -> std::filesystem::path {
+    std::filesystem::path root(
+        LR"(C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA)");
+    for (auto& file : std::filesystem::directory_iterator(root)) {
+      if (std::filesystem::exists(file.path() / "include" / "cuda.h")) {
+        return file.path();
+      }
+    }
+    return {};
+  }();
+#else
+  static std::filesystem::path cached_path = "/usr/local/cuda";
+#endif
+  return cached_path;
+}
+
 // Return the --include-path args used for invoking NVRTC.
 const std::vector<std::string>& include_path_args() {
  static std::vector<std::string> cached_args = []() {
@@ -47,26 +67,22 @@ const std::vector<std::string>& include_path_args() {
    // Add path to CUDA runtime headers, try local-installed python package
    // first and then system-installed headers.
    path = root_dir.parent_path() / "nvidia" / "cuda_runtime" / "include";
-    if (std::filesystem::exists(path)) {
-      args.push_back(fmt::format("--include-path={}", path.string()));
-    } else {
+    if (!std::filesystem::exists(path)) {
      const char* home = std::getenv("CUDA_HOME");
      if (!home) {
        home = std::getenv("CUDA_PATH");
      }
-#if defined(__linux__)
-      if (!home) {
-        home = "/usr/local/cuda";
+      path = home ? std::filesystem::path(home) : default_cuda_toolkit_path();
+      if (!path.empty()) {
+        path = path / "include";
      }
-#endif
-      if (home && std::filesystem::exists(home)) {
-        args.push_back(fmt::format("--include-path={}/include", home));
-      } else {
+      if (path.empty() || !std::filesystem::exists(path)) {
        throw std::runtime_error(
            "Can not find locations of CUDA headers, please set environment "
            "variable CUDA_HOME or CUDA_PATH.");
      }
    }
+    args.push_back(fmt::format("--include-path={}", path.string()));
    return args;
  }();
  return cached_args;
@@ -231,6 +247,7 @@ constexpr const char* g_include_names[] = {
    INCLUDE_PREFIX "config.h",
    INCLUDE_PREFIX "complex.cuh",
    INCLUDE_PREFIX "fp16_math.cuh",
+    INCLUDE_PREFIX "hadamard.cuh",
    INCLUDE_PREFIX "indexing.cuh",
    INCLUDE_PREFIX "scatter_ops.cuh",
    INCLUDE_PREFIX "unary_ops.cuh",
@@ -247,6 +264,7 @@ constexpr const char* g_headers[] = {
    jit_source_config,
    jit_source_complex,
    jit_source_fp16_math,
+    jit_source_hadamard,
    jit_source_indexing,
    jit_source_scatter_ops,
    jit_source_unary_ops,
@@ -422,20 +440,28 @@ CUfunction JitModule::get_kernel(
  return get_kernel_and_dims(kernel_name, std::move(configure_kernel)).first;
 }

-std::unordered_map<std::string, JitModule>& get_jit_module_cache() {
-  static std::unordered_map<std::string, JitModule> map;
-  return map;
-}
-
 JitModule& get_jit_module(
    const mlx::core::Device& device,
    const std::string& name,
    const KernelBuilder& builder,
-    bool cache) {
-  auto& map = get_jit_module_cache();
-  auto it = map.find(name);
-  if (it == map.end()) {
-    it = map.try_emplace(name, cu::device(device), name, builder, cache).first;
+    bool use_disk_cache) {
+  // The cache are leak intentionally as user code may still be running JIT
+  // compiled code after main thread teardown.
+  static auto* cache = new std::unordered_map<std::string, JitModule>;
+  static auto* mtx = new std::shared_mutex;
+
+  {
+    std::shared_lock rlock(*mtx);
+    if (auto it = cache->find(name); it != cache->end()) {
+      return it->second;
+    }
+  }
+
+  std::unique_lock wlock(*mtx);
+  auto it = cache->find(name);
+  if (it == cache->end()) {
+    auto& d = cu::device(device);
+    it = cache->try_emplace(name, d, name, builder, use_disk_cache).first;
  }
  return it->second;
 }
@@ -91,11 +91,12 @@ class JitModule {
      Device& device,
      const std::string& module_name,
      const KernelBuilder& builder,
-      bool cache);
+      bool use_disk_cache);
  ~JitModule();

  JitModule(const JitModule&) = delete;
  JitModule& operator=(const JitModule&) = delete;
+
  CUfunction get_kernel(
      const std::string& kernel_name,
      std::function<void(CUfunction)> configure_kernel = nullptr);
@@ -109,8 +110,6 @@ class JitModule {
      kernels_;
 };

-std::unordered_map<std::string, JitModule>& get_jit_module_cache();
-
 JitModule& get_jit_module(
    const mlx::core::Device& device,
    const std::string& name,
@@ -89,7 +89,8 @@ using cuda_type_t = typename CTypeToCudaType<T>::type;
 template <typename T>
 inline constexpr bool is_floating_v =
    cuda::std::is_same_v<T, float> || cuda::std::is_same_v<T, double> ||
-    cuda::std::is_same_v<T, float16_t> || cuda::std::is_same_v<T, bfloat16_t>;
+    cuda::std::is_same_v<T, float16_t> || cuda::std::is_same_v<T, bfloat16_t> ||
+    cuda::std::is_same_v<T, __half> || cuda::std::is_same_v<T, __nv_bfloat16>;

 // Type traits for detecting complex numbers.
 template <typename T>
@@ -279,7 +279,6 @@ void LayerNorm::eval_gpu(
          kernel,
          n_rows,
          block_dim(),
-          0,
          gpu_ptr<DataType>(x),
          gpu_ptr<DataType>(w),
          gpu_ptr<DataType>(b),
@@ -392,7 +391,6 @@ void LayerNormVJP::eval_gpu(
                kernel,
                n_rows,
                block_dim(),
-                0,
                gpu_ptr<DataType>(x),
                gpu_ptr<DataType>(w),
                gpu_ptr<DataType>(g),
@@ -150,7 +150,6 @@ void LogSumExp::eval_gpu(const std::vector<array>& inputs, array& out) {
          kernel,
          n_rows,
          block_dim(),
-          0,
          gpu_ptr<DataType>(in),
          gpu_ptr<DataType>(out),
          axis_size);
@@ -137,9 +137,15 @@ class LRUCache {

 // Turn a POD struct into a container key by doing bytes compare.
 //
+// IMPORTANT: Do not use aggregate init on the pod field (key.pod = {...}).
+// It creates a stack temporary whose padding bytes are uninitialized, and
+// trivial copy-assignment copies the entire struct including padding —
+// breaking the memcmp-based comparison. Set fields individually instead.
+//
 // Usage:
 //   BytesKey<MyKey> key;
-//   key.pod = { ... };
+//   key.pod.field1 = value1;
+//   key.pod.field2 = value2;
 template <typename T>
 struct BytesKey {
  T pod;
@@ -2,7 +2,9 @@

 #include "mlx/backend/common/matmul.h"
 #include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/gemms/block_mask.h"
 #include "mlx/backend/cuda/gemms/cublas_gemm.h"
+#include "mlx/backend/cuda/gemms/gather_gemm.h"
 #include "mlx/backend/cuda/gemms/gemv.h"
 #include "mlx/backend/cuda/gemms/grouped_gemm.h"
 #include "mlx/backend/gpu/copy.h"
@@ -136,40 +138,6 @@ void gemm_and_bias(
      encoder, out, a, b, batch_shape, a_batch_strides, b_batch_strides, alpha);
 }

-void gather_mm_rhs(
-    const array& a_,
-    const array& b_,
-    const array& indices_,
-    array& out,
-    cu::CommandEncoder& encoder,
-    Stream s) {
-  if (a_.size() / a_.shape(-2) / a_.shape(-1) != indices_.size()) {
-    throw std::runtime_error("[gather_mm] Broadcasting lhs is not supported.");
-  }
-
-  int group_count = b_.size() / b_.shape(-1) / b_.shape(-2);
-  if (group_count > 1024) {
-    throw std::runtime_error(
-        "[gather_mm] Group count can not be larger than 1024.");
-  }
-
-  auto [a_transposed, lda, a] = ensure_batch_contiguous(a_, encoder, s);
-  auto [b_transposed, ldb, b] = ensure_batch_contiguous(b_, encoder, s);
-  auto indices = ensure_row_contiguous(indices_, encoder, s);
-
-  cutlass_grouped_gemm_unaligned(
-      a_transposed,
-      lda,
-      b_transposed,
-      ldb,
-      group_count,
-      a,
-      b,
-      indices,
-      out,
-      encoder);
-}
-
 } // namespace

 void Matmul::eval_gpu(const std::vector<array>& inputs, array& out) {
@@ -203,6 +171,82 @@ void Matmul::eval_gpu(const std::vector<array>& inputs, array& out) {
      encoder, M, N, K, a_transposed, lda, b_transposed, ldb, out, a, b);
 }

+void BlockMaskedMM::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("BlockMaskedMM::eval_gpu");
+  if (!issubdtype(out.dtype(), floating)) {
+    throw std::runtime_error(
+        "[BlockMaskedMM] Does not yet support non-floating point types.");
+  }
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  auto& a_pre = inputs[0];
+  auto& b_pre = inputs[1];
+
+  // Return 0s if either input is empty.
+  if (a_pre.size() == 0 || b_pre.size() == 0) {
+    array zero(0, a_pre.dtype());
+    encoder.add_temporary(zero);
+    fill_gpu(zero, out, s);
+    return;
+  }
+
+  int M = a_pre.shape(-2);
+  int N = b_pre.shape(-1);
+  int K = a_pre.shape(-1);
+
+  if (M == 0 || N == 0) {
+    return;
+  }
+  if (K == 0) {
+    array zero(0, a_pre.dtype());
+    encoder.add_temporary(zero);
+    fill_gpu(zero, out, s);
+    return;
+  }
+
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+
+  bool has_op_mask = inputs.size() > 3;
+  bool has_out_mask = inputs.size() == 3 || inputs.size() == 5;
+
+  int64_t batch_count = out.size() / (int64_t(M) * N);
+
+  bool a_transposed;
+  int64_t lda;
+  array a = a_pre;
+  bool b_transposed;
+  int64_t ldb;
+  array b = b_pre;
+
+  if (has_op_mask) {
+    // Fused copy + mask in a single pass per matrix.
+    auto& lhs_mask = inputs[inputs.size() - 2];
+    auto& rhs_mask = inputs[inputs.size() - 1];
+    a = copy_with_block_mask(
+        encoder, a_pre, lhs_mask, block_size_, M, K, batch_count);
+    b = copy_with_block_mask(
+        encoder, b_pre, rhs_mask, block_size_, K, N, batch_count);
+    a_transposed = false;
+    lda = K;
+    b_transposed = false;
+    ldb = N;
+  } else {
+    std::tie(a_transposed, lda, a) = check_transpose(encoder, s, a_pre);
+    std::tie(b_transposed, ldb, b) = check_transpose(encoder, s, b_pre);
+  }
+
+  // Run GEMM.
+  gemm_and_bias(
+      encoder, M, N, K, a_transposed, lda, b_transposed, ldb, out, a, b);
+
+  // Apply output mask.
+  if (has_out_mask) {
+    auto& out_mask = inputs[2];
+    apply_block_mask(encoder, out, out_mask, block_size_, M, N, batch_count);
+  }
+}
+
 void AddMM::eval_gpu(const std::vector<array>& inputs, array& out) {
  nvtx3::scoped_range r("AddMM::eval_gpu");
  auto& s = stream();
@@ -327,14 +371,12 @@ void GatherMM::eval_gpu(const std::vector<array>& inputs, array& out) {
  auto& encoder = cu::get_command_encoder(s);

  assert(inputs.size() == 4);
-  auto& a = inputs[0];
-  auto& b = inputs[1];
-  auto& lhs_indices = inputs[2];
-  auto& rhs_indices = inputs[3];
+  auto& a_pre = inputs[0];
+  auto& b_pre = inputs[1];

  // Return 0s if either input is empty.
-  if (a.size() == 0 || b.size() == 0) {
-    array zero(0, a.dtype());
+  if (a_pre.size() == 0 || b_pre.size() == 0) {
+    array zero(0, a_pre.dtype());
    encoder.add_temporary(zero);
    fill_gpu(zero, out, s);
    return;
@@ -343,31 +385,87 @@ void GatherMM::eval_gpu(const std::vector<array>& inputs, array& out) {
  out.set_data(cu::malloc_async(out.nbytes(), encoder));

  // Extract shapes from inputs.
-  int M = a.shape(-2);
-  int N = b.shape(-1);
-  int K = a.shape(-1);
+  int M = a_pre.shape(-2);
+  int N = b_pre.shape(-1);
+  int K = a_pre.shape(-1);
+
+  auto [a_transposed, lda, a] = ensure_batch_contiguous(a_pre, encoder, s);
+  auto [b_transposed, ldb, b] = ensure_batch_contiguous(b_pre, encoder, s);
+  auto lhs_indices = ensure_row_contiguous(inputs[2], encoder, s);
+  auto rhs_indices = ensure_row_contiguous(inputs[3], encoder, s);

  // We are walking a in order and b is also in order so we can batch up the
  // matmuls and reuse reading a and b.
  if (M == 1 && right_sorted_ == true) {
-    gather_mm_rhs(a, b, rhs_indices, out, encoder, s);
+    cutlass_grouped_gemm_unaligned(
+        a_transposed,
+        lda,
+        b_transposed,
+        ldb,
+        b.size() / b.shape(-1) / b.shape(-2), // group_count
+        a,
+        b,
+        rhs_indices,
+        out,
+        encoder);
    return;
  }

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

-  throw std::runtime_error("NYI");
+  cutlass_gather_mm(
+      a_transposed, b_transposed, a, b, lhs_indices, rhs_indices, out, encoder);
+}
+
+void SegmentedMM::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("SegmentedMM::eval_gpu");
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  assert(inputs.size() == 3);
+  auto& a_pre = inputs[0];
+  auto& b_pre = inputs[1];
+  auto& segments_pre = inputs[2];
+
+  // Return zeros if output is empty or either input is empty.
+  if (out.size() == 0 || a_pre.size() == 0 || b_pre.size() == 0) {
+    array zero(0, a_pre.dtype());
+    encoder.add_temporary(zero);
+    fill_gpu(zero, out, s);
+    return;
+  }
+
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+
+  int M = a_pre.shape(-2);
+  int N = b_pre.shape(-1);
+  int num_segments = segments_pre.size() / 2;
+
+  auto [a_transposed, lda, a] = check_transpose(encoder, s, a_pre);
+  auto [b_transposed, ldb, b] = check_transpose(encoder, s, b_pre);
+  auto segments = ensure_row_contiguous(segments_pre, encoder, s);
+
+  cutlass_segmented_mm(
+      a_transposed,
+      lda,
+      b_transposed,
+      ldb,
+      num_segments,
+      M,
+      N,
+      a,
+      b,
+      segments,
+      out,
+      encoder);
 }

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

-NO_GPU(BlockMaskedMM)
-NO_GPU(FFT)
-NO_GPU(GatherQMM)
-NO_GPU(Hadamard)
 NO_GPU_MULTI(LUF)
 NO_GPU_MULTI(QRF)
-NO_GPU(SegmentedMM)
 NO_GPU_MULTI(SVD)
 NO_GPU(Inverse)
 NO_GPU(Cholesky)
 NO_GPU_MULTI(Eig)
 NO_GPU_MULTI(Eigh)
-NO_GPU(MaskedScatter)

 namespace distributed {
 NO_GPU_MULTI(Send)
@@ -1,8 +1,8 @@
 // Copyright © 2025 Apple Inc.

+#include "mlx/backend/common/quantized.h"
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/kernel_utils.cuh"
-#include "mlx/backend/cuda/quantized/quantized_utils.cuh"
 #include "mlx/dtype_utils.h"

 #include <cooperative_groups.h>
@@ -27,8 +27,8 @@ affine_quantize(const T* w, uint8_t* out, T* scales, T* biases, size_t size) {
  constexpr float eps = 1e-7;
  constexpr int simd_size = WARP_SIZE;
  constexpr float n_bins = (1 << bits) - 1;
-  constexpr int pack_factor = get_pack_factor<bits, 8>();
-  constexpr int bytes_per_pack = get_bytes_per_pack<bits>();
+  constexpr int pack_factor = get_pack_factor(bits, 8);
+  constexpr int bytes_per_pack = get_bytes_per_pack(bits);
  constexpr int values_per_reduce = group_size / simd_size;
  constexpr int writes_per_reduce = pack_factor / values_per_reduce;
  constexpr int writes_per_pack =
@@ -142,8 +142,8 @@ __global__ void affine_dequantize(

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

-  constexpr int pack_factor = get_pack_factor<bits, 8>();
-  constexpr int bytes_per_pack = get_bytes_per_pack<bits>();
+  constexpr int pack_factor = get_pack_factor(bits, 8);
+  constexpr int bytes_per_pack = get_bytes_per_pack(bits);

  size_t offset = tidx + grid_dim_x * size_t(tidy);
  size_t oindex = offset * pack_factor;
@@ -226,6 +226,45 @@ __global__ void affine_dequantize(

 } // namespace cu

+template <typename F>
+void dispatch_groups(int group_size, F&& f) {
+  switch (group_size) {
+    case 32:
+      f(std::integral_constant<int, 32>{});
+      break;
+    case 64:
+      f(std::integral_constant<int, 64>{});
+      break;
+    case 128:
+      f(std::integral_constant<int, 128>{});
+      break;
+  }
+}
+
+template <typename F>
+void dispatch_bits(int bits, F&& f) {
+  switch (bits) {
+    case 2:
+      f(std::integral_constant<int, 2>{});
+      break;
+    case 3:
+      f(std::integral_constant<int, 3>{});
+      break;
+    case 4:
+      f(std::integral_constant<int, 4>{});
+      break;
+    case 5:
+      f(std::integral_constant<int, 5>{});
+      break;
+    case 6:
+      f(std::integral_constant<int, 6>{});
+      break;
+    case 8:
+      f(std::integral_constant<int, 8>{});
+      break;
+  }
+}
+
 void affine_quantize(
    const array& w,
    array& wq,
@@ -259,7 +298,6 @@ void affine_quantize(
            kernel,
            num_blocks,
            block_dims,
-            0,
            gpu_ptr<T>(w),
            gpu_ptr<uint8_t>(wq),
            gpu_ptr<T>(scales),
@@ -315,7 +353,6 @@ void affine_dequantize(
            kernel,
            num_blocks,
            block_dims,
-            0,
            gpu_ptr<uint8_t>(wq),
            gpu_ptr<T>(scales),
            gpu_ptr<T>(biases),
@@ -13,39 +13,26 @@ namespace mlx::core {

 namespace {

-// Currently cublas supports only mxfp8 and nvfp4
-// quantization modes for block scaled quantization
-cudaDataType_t qmode_to_cublas_scale_dtype(std::string mode) {
-  if (mode == "mxfp8") {
-    return CUDA_R_8F_UE8M0;
-  } else if (mode == "nvfp4") {
-    return CUDA_R_8F_UE4M3;
-  } else {
-    throw std::runtime_error(
-        fmt::format("Unsupported quantization mode in CublasQQMM: {}.", mode));
-  }
-}
+struct QuantModeConfig {
+  cudaDataType_t data_type;
+  cudaDataType_t scale_dtype;
+  cublasLtMatmulMatrixScale_t scale_mode;
+};

-cudaDataType_t qmode_to_cublas_dtype(std::string mode) {
+QuantModeConfig get_quant_mode_config(const std::string& mode) {
  if (mode == "mxfp8") {
-    return CUDA_R_8F_E4M3;
+    return {
+        CUDA_R_8F_E4M3,
+        CUDA_R_8F_UE8M0,
+        CUBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0};
  } else if (mode == "nvfp4") {
-    return CUDA_R_4F_E2M1;
-  } else {
-    throw std::runtime_error(
-        fmt::format("Unsupported quantization mode in CublasQQMM: {}.", mode));
-  }
-}
-
-cublasLtMatmulMatrixScale_t qmode_to_cublas_scale_mode(std::string mode) {
-  if (mode == "mxfp8") {
-    return CUBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0;
-  } else if (mode == "nvfp4") {
-    return CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3;
-  } else {
-    throw std::runtime_error(
-        fmt::format("Unsupported quantization mode in CublasQQMM: {}.", mode));
+    return {
+        CUDA_R_4F_E2M1,
+        CUDA_R_8F_UE4M3,
+        CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3};
  }
+  throw std::runtime_error(
+      fmt::format("Unsupported quantization mode in CublasQQMM: {}.", mode));
 }

 } // namespace
@@ -64,21 +51,21 @@ CublasQQMM::CublasQQMM(
    int64_t a_batch_stride,
    int64_t b_batch_stride,
    Dtype out_dtype,
-    std::string qmode) {
+    const std::string& qmode) {
+  auto config = get_quant_mode_config(qmode);
+
  // The compute type must be CUBLAS_COMPUTE_32F.
  // The scale type must be CUDA_R_32F.
  cudaDataType_t scale_type = CUDA_R_32F;
  cublasComputeType_t gemm_compute_type = CUBLAS_COMPUTE_32F;
  cudaDataType_t output_type =
      cublas_utils::dtype_to_cublas_type(out_dtype, "CublasQQMM");
-  cudaDataType_t data_type = qmode_to_cublas_dtype(qmode);
-  quantization_mode_ = std::string(qmode);

  init_base(
      device,
      scale_type,
      gemm_compute_type,
-      data_type,
+      config.data_type,
      output_type,
      a_transposed,
      a_rows,
@@ -92,8 +79,8 @@ CublasQQMM::CublasQQMM(
      a_batch_stride,
      b_batch_stride);

-  a_scale_mode_ = qmode_to_cublas_scale_mode(qmode);
-  b_scale_mode_ = qmode_to_cublas_scale_mode(qmode);
+  a_scale_mode_ = config.scale_mode;
+  b_scale_mode_ = config.scale_mode;

  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
      matmul_desc_,
@@ -123,7 +110,7 @@ CublasQQMM::CublasQQMM(
    int64_t b_batch_stride,
    int64_t c_batch_stride,
    Dtype out_dtype,
-    std::string qmode)
+    const std::string& qmode)
    : CublasQQMM(
          device,
          a_transposed,
@@ -158,7 +145,35 @@ void CublasQQMM::run(
    const array& b,
    const array& a_scale,
    const array& b_scale,
-    float alpha) {
+    const array& alpha,
+    const array& beta) {
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(a_scale);
+  encoder.set_input_array(b_scale);
+  encoder.set_input_array(alpha);
+  encoder.set_input_array(beta);
+  encoder.set_output_array(out);
+
+  execute(
+      encoder,
+      gpu_ptr<void>(out),
+      gpu_ptr<void>(a),
+      gpu_ptr<void>(b),
+      gpu_ptr<void>(a_scale),
+      gpu_ptr<void>(b_scale),
+      nullptr,
+      gpu_ptr<void>(alpha),
+      gpu_ptr<void>(beta));
+}
+
+void CublasQQMM::run(
+    cu::CommandEncoder& encoder,
+    array& out,
+    const array& a,
+    const array& b,
+    const array& a_scale,
+    const array& b_scale) {
  encoder.set_input_array(a);
  encoder.set_input_array(b);
  encoder.set_input_array(a_scale);
@@ -172,20 +187,13 @@ void CublasQQMM::run(
      gpu_ptr<void>(b),
      gpu_ptr<void>(a_scale),
      gpu_ptr<void>(b_scale),
-      nullptr,
-      alpha);
+      nullptr);
 }

-void CublasQQMM::execute(
+void CublasQQMM::set_scales_ptrs(
    cu::CommandEncoder& encoder,
-    void* out,
-    const void* a,
-    const void* b,
    const void* a_scale,
-    const void* b_scale,
-    const void* c,
-    float alpha /* = 1 */,
-    float beta /* = 0 */) {
+    const void* b_scale) {
  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
      matmul_desc_,
      CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,
@@ -196,6 +204,49 @@ void CublasQQMM::execute(
      CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
      &a_scale,
      sizeof(a_scale)));
+}
+
+void CublasQQMM::execute(
+    cu::CommandEncoder& encoder,
+    void* out,
+    const void* a,
+    const void* b,
+    const void* a_scale,
+    const void* b_scale,
+    const void* c,
+    const void* alpha,
+    const void* beta) {
+  set_scales_ptrs(encoder, a_scale, b_scale);
+  // alpha and beta are both should be device pointers for nvfp4
+  // by default cublas uses host pointers
+  // https://docs.nvidia.com/cuda/cublas/#cublasltpointermode-t
+  cublasLtPointerMode_t pointer_mode = CUBLASLT_POINTER_MODE_DEVICE;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_POINTER_MODE,
+      &pointer_mode,
+      sizeof(pointer_mode)));
+  execute_matmul(encoder, out, a, b, c, alpha, beta);
+}
+
+void CublasQQMM::execute(
+    cu::CommandEncoder& encoder,
+    void* out,
+    const void* a,
+    const void* b,
+    const void* a_scale,
+    const void* b_scale,
+    const void* c,
+    const float alpha /* = 1 */,
+    const float beta /* = 0 */) {
+  set_scales_ptrs(encoder, a_scale, b_scale);
+  // alpha and beta are both should be 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)));

  const void* alpha_ptr = &alpha;
  const void* beta_ptr = &beta;
@@ -25,7 +25,7 @@ class CublasQQMM : public CublasMatmulBase {
      int64_t a_batch_stride,
      int64_t b_batch_stride,
      Dtype out_dtype,
-      std::string quantization_mode);
+      const std::string& quantization_mode);

  CublasQQMM(
      cu::Device& device,
@@ -43,7 +43,7 @@ class CublasQQMM : public CublasMatmulBase {
      int64_t b_batch_stride,
      int64_t c_batch_stride,
      Dtype out_dtype,
-      std::string quantization_mode);
+      const std::string& quantization_mode);

  void run(
      cu::CommandEncoder& encoder,
@@ -52,20 +52,22 @@ class CublasQQMM : public CublasMatmulBase {
      const array& b,
      const array& a_scale,
      const array& b_scale,
-      float alpha = 1.0f);
+      const array& alpha,
+      const array& beta);

- private:
-  void run_batched(
+  void run(
      cu::CommandEncoder& encoder,
      array& out,
      const array& a,
      const array& b,
      const array& a_scale,
-      const array& b_scale,
-      const Shape& batch_shape,
-      const Strides& a_batch_strides,
-      const Strides& b_batch_strides,
-      float alpha);
+      const array& b_scale);
+
+ private:
+  void set_scales_ptrs(
+      cu::CommandEncoder& encoder,
+      const void* a_scale,
+      const void* b_scale);

  void execute(
      cu::CommandEncoder& encoder,
@@ -75,10 +77,20 @@ class CublasQQMM : public CublasMatmulBase {
      const void* a_scale,
      const void* b_scale,
      const void* c,
-      float alpha = 1,
-      float beta = 0);
+      const void* alpha,
+      const void* beta);
+
+  void execute(
+      cu::CommandEncoder& encoder,
+      void* out,
+      const void* a,
+      const void* b,
+      const void* a_scale,
+      const void* b_scale,
+      const void* c,
+      const float alpha = 1.0f,
+      const float beta = 0.0f);

-  std::string quantization_mode_;
  cublasLtMatmulMatrixScale_t a_scale_mode_;
  cublasLtMatmulMatrixScale_t b_scale_mode_;
  cublasLtMatmulMatrixScale_t c_scale_mode_;
--- 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