fix(ci): use branch_to_build input for checkout and artifact naming

Builds mlx wheels for Python 3.11 and 3.12 on macos-14 (arm64)
on push to main, metal-*, q-*, attn-mask-fix, fix-rope branches
or via workflow_dispatch. Wheels uploaded as artifacts (30d retention).
Tagged commits (v*) also publish a GitHub Release.

.github/actions/build-cuda-release/action.yml

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

.github/actions/build-linux-release/action.yml

@@ -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 \

.github/actions/build-macos/action.yml

@@ -12,12 +12,12 @@ runs:
       run: |
         pip install --upgrade pip
         pip install cmake setuptools typing_extensions
-        pip install -e . -v
+        pip install -e ".[dev]" -v
 
     - name: Install tests dependencies
       shell: bash -l {0}
       run: |
-        pip install numpy torch tensorflow
+        pip install tensorflow
 
     - name: Run Python tests
       shell: bash -l {0}
@@ -45,15 +45,17 @@ runs:
         cd build
         cmake ..
         make -j $(sysctl -n hw.ncpu)
-    
+
     - name: Run CPP tests
       shell: bash -l {0}
       env:
         DEVICE: gpu
         METAL_DEVICE_WRAPPER_TYPE: 1
         METAL_DEBUG_ERROR_MODE: 0
-      run: ./build/tests/tests
-    
+      run: |
+        ./build/tests/tests
+        ./build/tests/test_teardown
+
     - name: Build small binary with JIT
       shell: bash -l {0}
       run: |

.github/actions/setup-linux/action.yml

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

.github/actions/test-linux/action.yml

@@ -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::"

.github/actions/test-windows/action.yml

@@ -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::"

.github/workflows/build-wheels.yml

@@ -0,0 +1,94 @@
+name: Build macOS arm64 wheels
+
+on:
+  push:
+    branches:
+      - main
+      - 'metal-*'
+      - 'q-*'
+      - attn-mask-fix
+      - fix-rope
+  workflow_dispatch:
+    inputs:
+      branch_to_build:
+        description: 'Branch to build (optional, defaults to current ref)'
+        required: false
+        default: ''
+
+concurrency:
+  group: build-${{ github.ref }}-${{ github.event.inputs.branch_to_build }}
+  cancel-in-progress: true
+
+jobs:
+  build:
+    name: Build wheel (Python ${{ matrix.python }})
+    runs-on: macos-14
+    timeout-minutes: 60
+    strategy:
+      fail-fast: false
+      matrix:
+        python: ['3.11', '3.12']
+    env:
+      CMAKE_BUILD_PARALLEL_LEVEL: '4'
+    steps:
+      - name: Determine target branch
+        id: branch
+        run: |
+          NAME="${{ github.event.inputs.branch_to_build }}"
+          if [ -z "$NAME" ]; then
+            NAME="${{ github.ref_name }}"
+          fi
+          # Sanitize for artifact naming (replace / with -)
+          SAFE_NAME=$(echo "$NAME" | tr '/' '-')
+          echo "name=$NAME" >> $GITHUB_OUTPUT
+          echo "safe_name=$SAFE_NAME" >> $GITHUB_OUTPUT
+          echo "Target branch: $NAME (safe: $SAFE_NAME)"
+
+      - name: Checkout
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ steps.branch.outputs.name }}
+          submodules: recursive
+
+      - name: Set up Python ${{ matrix.python }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python }}
+
+      - name: Cache pip
+        uses: actions/cache@v4
+        with:
+          path: |
+            ~/.cache/pip
+            ~/Library/Caches/pip
+          key: pip-${{ runner.os }}-py${{ matrix.python }}-${{ hashFiles('CMakeLists.txt', 'setup.py', 'pyproject.toml') }}
+          restore-keys: |
+            pip-${{ runner.os }}-py${{ matrix.python }}-
+
+      - name: Install build dependencies
+        run: |
+          python -m pip install -U pip wheel build setuptools cmake nanobind
+
+      - name: Build wheel
+        run: |
+          mkdir -p ./wheels
+          pip wheel --no-deps . -w ./wheels
+
+      - name: List built wheels
+        run: ls -lh ./wheels
+
+      - name: Upload wheel artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: mlx-${{ steps.branch.outputs.safe_name }}-py${{ matrix.python }}-wheels
+          path: ./wheels/*.whl
+          retention-days: 30
+          if-no-files-found: error
+
+      - name: Create GitHub Release (on tag)
+        if: startsWith(github.ref, 'refs/tags/v')
+        uses: softprops/action-gh-release@v2
+        with:
+          files: ./wheels/*.whl
+          fail_on_unmatched_files: false
+          generate_release_notes: true

.github/workflows/documentation.yml

@@ -25,4 +25,4 @@ jobs:
     steps:
       - name: Deploy to GitHub Pages
         id: deployment
-        uses: actions/deploy-pages@v4
+        uses: actions/deploy-pages@v5

.github/workflows/nightly.yml

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

.github/workflows/release.yml

@@ -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 }}
@@ -116,7 +116,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 +124,7 @@ jobs:
           if-no-files-found: error
       - name: Upload Metal artifacts
         if: matrix.python-version == '3.10'
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
         with:
           overwrite: true
           name: mlx-metal
@@ -146,13 +146,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 +169,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 +197,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 +220,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 +243,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

CMakeLists.txt

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

benchmarks/python/block_masked_mm_bench.py

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

benchmarks/python/masked_scatter.py

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

benchmarks/python/sdpa_bench.py

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

benchmarks/python/segmented_mm_bench.py

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

benchmarks/python/slice_update_bench.py

@@ -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}"
+            )

docs/logo/mlx_logo.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(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>

docs/logo/mlx_logo_dark.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>

docs/src/dev/extensions.rst

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

docs/src/index.rst

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

docs/src/install.rst

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

docs/src/python/devices_and_streams.rst

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

docs/src/python/fft.rst

@@ -20,5 +20,7 @@ FFT
   irfft2
   rfftn
   irfftn
+  fftfreq
+  rfftfreq
   fftshift
   ifftshift

docs/src/python/linalg.rst

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

docs/src/python/nn.rst

@@ -175,6 +175,7 @@ In detail:
    value_and_grad
    quantize
    average_gradients
+   fsdp_apply_gradients
 
 .. toctree::
 

docs/src/python/printoptions.rst

@@ -0,0 +1,12 @@
+Print Options
+===============
+
+.. currentmodule:: mlx.core
+
+.. autosummary::
+  :toctree: _autosummary
+
+  PrintOptions
+  set_printoptions
+  printoptions
+  get_printoptions

docs/src/usage/numpy.rst

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

examples/extensions/axpby/axpby.cpp

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

examples/extensions/pyproject.toml

@@ -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"

examples/extensions/requirements.txt

@@ -1,4 +1,4 @@
 setuptools>=42
 cmake>=3.25
 mlx>=0.21.0
-nanobind==2.10.2
+nanobind==2.12.0

mlx/CMakeLists.txt

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

mlx/backend/common/common.cpp

@@ -19,27 +19,28 @@ void AsStrided::eval(const std::vector<array>& inputs, array& out) {
         "AsStrided must be used with row contiguous arrays only.");
   }
 
-  // Compute the flags given the shape and strides
-  bool row_contiguous = true, col_contiguous = true;
-  size_t r = 1, c = 1;
-  for (int i = strides_.size() - 1, j = 0; i >= 0; i--, j++) {
-    row_contiguous &= (r == strides_[i]) || (shape_[i] == 1);
-    col_contiguous &= (c == strides_[j]) || (shape_[j] == 1);
-    r *= shape_[i];
-    c *= shape_[j];
+  auto [no_bsx_size, row_contiguous, col_contiguous] =
+      check_contiguity(shape_, strides_);
+
+  int64_t l = 0, h = 0;
+  bool has_negative_stride = false;
+  for (int i = 0; i < strides_.size(); i++) {
+    auto delta = strides_[i] * (shape_[i] - 1);
+    if (strides_[i] >= 0) {
+      h += delta;
+    } else {
+      l += delta;
+      has_negative_stride |= shape_[i] > 1;
+    }
   }
+  size_t data_size = out.size() == 0 ? 0 : (h - l) + 1;
+
   auto flags = in.flags();
-  // TODO: Compute the contiguous flag in a better way cause now we are
-  //       unnecessarily strict.
-  flags.contiguous = row_contiguous || col_contiguous;
+  flags.contiguous =
+      out.size() == 0 || (!has_negative_stride && no_bsx_size == data_size);
   flags.row_contiguous = row_contiguous;
   flags.col_contiguous = col_contiguous;
 
-  // There is no easy way to compute the actual data size so we use out.size().
-  // The contiguous flag will almost certainly not be set so no code should
-  // rely on data_size anyway.
-  size_t data_size = out.size();
-
   return out.copy_shared_buffer(in, strides_, flags, data_size, offset_);
 }
 

mlx/backend/common/quantized.h

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

mlx/backend/common/utils.h

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

mlx/backend/cpu/compiled.cpp

@@ -10,7 +10,6 @@
 #include <fmt/format.h>
 
 #include "mlx/backend/common/compiled.h"
-#include "mlx/backend/cpu/compiled_preamble.h"
 #include "mlx/backend/cpu/encoder.h"
 #include "mlx/backend/cpu/jit_compiler.h"
 #include "mlx/device.h"
@@ -316,7 +315,9 @@ void Compiled::eval_cpu(
   // Get the function
   auto fn_ptr = compile(kernel_name, [&, contiguous = contiguous]() {
     std::ostringstream kernel;
-    kernel << get_kernel_preamble() << std::endl;
+    kernel << std::get<2>(JitCompiler::get_preamble()) << std::endl;
+    kernel << "using namespace mlx::core;" << std::endl;
+    kernel << "using namespace mlx::core::detail;" << std::endl;
     kernel << "extern \"C\"  {" << std::endl;
     build_kernel(
         kernel,

mlx/backend/cpu/compiled_preamble.h

@@ -9,4 +9,4 @@
 #include "mlx/backend/cpu/binary_ops.h"
 // clang-format on
 
-const char* get_kernel_preamble();
+const char* get_prebuilt_preamble();

mlx/backend/cpu/indexing.cpp

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

mlx/backend/cpu/jit_compiler.cpp

@@ -1,6 +1,8 @@
 // Copyright © 2024 Apple Inc.
 
 #include "mlx/backend/cpu/jit_compiler.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cpu/compiled_preamble.h"
 
 #include <algorithm>
 #include <sstream>
@@ -86,30 +88,61 @@ const VisualStudioInfo& GetVisualStudioInfo() {
 
 #endif // _MSC_VER
 
+const std::tuple<bool, std::string, std::string>& JitCompiler::get_preamble() {
+  static auto preamble = []() -> std::tuple<bool, std::string, std::string> {
+    // Check whether the headers are shipped with the binary, if so use the
+    // preamble from the headers, otherwise use the prebuilt one embeded in
+    // binary, which may not work with all compilers.
+    auto root_dir = current_binary_dir();
+#if !defined(_WIN32)
+    root_dir = root_dir.parent_path();
+#endif
+    auto include_dir = root_dir / "include";
+    if (std::filesystem::exists(include_dir / "mlx")) {
+      return std::make_tuple(
+          true,
+          include_dir.string(),
+          "#include \"mlx/backend/cpu/compiled_preamble.h\"\n");
+    } else {
+      return std::make_tuple(false, "", get_prebuilt_preamble());
+    }
+  }();
+  return preamble;
+}
+
 std::string JitCompiler::build_command(
     const std::filesystem::path& dir,
     const std::string& source_file_name,
     const std::string& shared_lib_name) {
+  auto& [use_include, include_dir, preamble] = get_preamble();
 #ifdef _MSC_VER
+  std::string extra_flags;
+  if (use_include) {
+    extra_flags += fmt::format("/I \"{}\"", include_dir);
+  }
   const VisualStudioInfo& info = GetVisualStudioInfo();
-  std::string libpaths;
   for (const std::string& lib : info.libpaths) {
-    libpaths += fmt::format(" /libpath:\"{0}\"", lib);
+    extra_flags += fmt::format(" /libpath:\"{}\"", lib);
   }
   return fmt::format(
       "\""
-      "cd /D \"{0}\" && "
-      "\"{1}\" /LD /EHsc /MD /Ox /nologo /std:c++17 \"{2}\" "
-      "/link /out:\"{3}\" {4} 2>&1"
+      "cd /D \"{}\" && "
+      "\"{}\" /LD /EHsc /MD /Ox /nologo /std:c++17 {} \"{}\" "
+      "/link /out:\"{}\" 2>&1"
       "\"",
       dir.string(),
       info.cl_exe,
+      extra_flags,
       source_file_name,
-      shared_lib_name,
-      libpaths);
+      shared_lib_name);
 #else
+  std::string extra_flags;
+  if (use_include) {
+    extra_flags = fmt::format("-I \"{}\"", include_dir);
+  }
   return fmt::format(
-      "g++ -std=c++17 -O3 -Wall -fPIC -shared \"{0}\" -o \"{1}\" 2>&1",
+      "g++ -std=c++17 -O3 -Wall -fPIC -shared {} \"{}\" -o \"{}\" 2>&1",
+      extra_flags,
       (dir / source_file_name).string(),
       (dir / shared_lib_name).string());
 #endif

mlx/backend/cpu/jit_compiler.h

@@ -7,6 +7,9 @@ namespace mlx::core {
 
 class JitCompiler {
  public:
+  // Return the includes that should be prepended to the source code.
+  static const std::tuple<bool, std::string, std::string>& get_preamble();
+
   // Build a shell command that compiles a source code file to a shared library.
   static std::string build_command(
       const std::filesystem::path& dir,

mlx/backend/cpu/luf.cpp

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

mlx/backend/cpu/make_compiled_preamble.ps1

@@ -15,13 +15,6 @@ $CONTENT = $CONTENT | Where-Object { $_.Trim() -ne '' }
 # Concatenate to string.
 $CONTENT = $CONTENT -join "`n"
 
-# Append extra content.
-$CONTENT = @"
-$($CONTENT)
-using namespace mlx::core;
-using namespace mlx::core::detail;
-"@
-
 # Convert each char to ASCII code.
 # Unlike the unix script that outputs string literal directly, the output from
 # MSVC is way too large to be embedded as string and compilation will fail, so
@@ -29,7 +22,7 @@ using namespace mlx::core::detail;
 $CHARCODES = ([System.Text.Encoding]::ASCII.GetBytes($CONTENT) -join ', ') + ', 0'
 
 $OUTPUT = @"
-const char* get_kernel_preamble() {
+const char* get_prebuilt_preamble() {
   static char preamble[] = { $CHARCODES };
   return preamble;
 }

mlx/backend/cpu/make_compiled_preamble.sh

@@ -30,12 +30,10 @@ fi
 CONTENT=$($GCC $CC_FLAGS -I "$SRCDIR" -E -P "$SRCDIR/mlx/backend/cpu/compiled_preamble.h" 2>/dev/null)
 
 cat << EOF > "$OUTPUT_FILE"
-const char* get_kernel_preamble() {
+const char* get_prebuilt_preamble() {
 return R"preamble(
 $INCLUDES
 $CONTENT
-using namespace mlx::core;
-using namespace mlx::core::detail;
 )preamble";
 }
 EOF

mlx/backend/cpu/primitives.cpp

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

mlx/backend/cpu/quantized.cpp

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

mlx/backend/cpu/sort.cpp

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

mlx/backend/cuda/CMakeLists.txt

@@ -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
@@ -117,9 +120,15 @@ target_compile_options(mlx
 target_compile_options(
   mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>")
 
-# Required for generating optimized CUTLASS code.
+# Ignore warnings from CUTLASS.
 target_compile_options(
-  mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-fno-strict-aliasing>")
+  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(
@@ -158,6 +167,7 @@ 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)
@@ -243,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)
 
@@ -266,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)

mlx/backend/cuda/binary/binary.cuh

@@ -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) {
@@ -331,7 +349,12 @@ 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,

mlx/backend/cuda/conv.cpp

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

mlx/backend/cuda/cublas_utils.cpp

@@ -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,8 +120,7 @@ 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(

mlx/backend/cuda/cublas_utils.h

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

mlx/backend/cuda/cuda_utils.h

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

mlx/backend/cuda/cudnn_utils.cpp

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

mlx/backend/cuda/cudnn_utils.h

@@ -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) {
@@ -182,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

mlx/backend/cuda/device.cpp

@@ -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"
@@ -55,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()) {
@@ -236,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) {
@@ -461,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_);
@@ -536,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;
 }
@@ -551,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) {
@@ -570,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

mlx/backend/cuda/device.h

@@ -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;
@@ -92,6 +95,10 @@ class CommandEncoder {
       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());
@@ -132,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_;
@@ -163,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_;
   }
@@ -198,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

mlx/backend/cuda/device/hadamard.cuh

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

mlx/backend/cuda/device/scatter.cuh

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

mlx/backend/cuda/device/slice_update.cuh

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

mlx/backend/cuda/eval.cpp

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

mlx/backend/cuda/fft.cu

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

mlx/backend/cuda/gemms/block_mask.cu

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

mlx/backend/cuda/gemms/block_mask.h

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

mlx/backend/cuda/gemms/gather_gemm.cu

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

mlx/backend/cuda/gemms/gather_gemm.h

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

mlx/backend/cuda/gemms/gemv.cu

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

mlx/backend/cuda/gemms/grouped_gemm.h

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

mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu

@@ -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 {
@@ -201,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,
@@ -281,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);
 
@@ -357,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

mlx/backend/cuda/hadamard.cu

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

mlx/backend/cuda/indexing.cpp

@@ -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,
@@ -435,4 +440,246 @@ void ScatterAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
       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

mlx/backend/cuda/jit_module.cpp

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

mlx/backend/cuda/jit_module.h

@@ -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,

mlx/backend/cuda/kernel_utils.cuh

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

mlx/backend/cuda/lru_cache.h

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

mlx/backend/cuda/matmul.cpp

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

mlx/backend/cuda/primitives.cpp

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

mlx/backend/cuda/quantized/affine_quantize.cu

@@ -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,

mlx/backend/cuda/quantized/cuda_fp4.h

@@ -1,100 +0,0 @@
-#pragma once
-
-struct __nv_fp8_e8m0 {
-  __device__ __nv_fp8_e8m0(float x) {
-    if (!std::isfinite(x)) {
-      __x = 0xFF;
-      return;
-    }
-    if (x < 0.0f) {
-      __x = 0x00;
-      return;
-    }
-    float le = std::log2f(x);
-    int n = static_cast<int>(std::nearbyintf(le));
-
-    n = n < -127 ? -127 : n;
-    n = n > 127 ? 127 : n;
-    __x = static_cast<uint8_t>(n + 127);
-  }
-
-  __device__ operator float() {
-    if (__x == 0xFF) {
-      return std::numeric_limits<float>::quiet_NaN();
-    }
-    return std::ldexp(1.0f, static_cast<int>(__x) - 127);
-  }
-
-  uint8_t __x{0};
-};
-
-struct __nv_fp4_e2m1 {
-  __device__ __nv_fp4_e2m1(float x) {
-    if (std::isnan(x)) {
-      __x = 0x7;
-      return;
-    }
-
-    const uint8_t sign_bit = (std::signbit(x)) ? 0x8 : 0x0;
-    x = std::abs(x);
-
-    if (x > 5.0f) {
-      __x = 0x7;
-    } else if (x >= 3.5f) {
-      __x = 0x6;
-    } else if (x > 2.5f) {
-      __x = 0x5;
-    } else if (x >= 1.75f) {
-      __x = 0x4;
-    } else if (x > 1.25f) {
-      __x = 0x3;
-    } else if (x >= 0.75f) {
-      __x = 0x2;
-    } else if (x > 0.25f) {
-      __x = 0x1;
-    } else {
-      __x = 0x0;
-    }
-    __x |= sign_bit;
-  }
-
-  __device__ operator float() {
-    static const float LUT[16] = {
-        0.0f,
-        0.5f,
-        1.0f,
-        1.5f,
-        2.0f,
-        3.0f,
-        4.0f,
-        6.0f,
-        -0.0f,
-        -0.5f,
-        -1.0f,
-        -1.5f,
-        -2.0f,
-        -3.0f,
-        -4.0f,
-        -6.0f};
-
-    return LUT[__x];
-  }
-  uint8_t __x{0};
-};
-
-struct __nv_fp4x4_e2m1 {
-  __device__ operator float4() {
-    float4 out;
-    auto bits = __high & 0xf;
-    out.x = float(*(__nv_fp4_e2m1*)(&bits));
-    bits = (__high >> 4) & 0xf;
-    out.y = float(*(__nv_fp4_e2m1*)(&bits));
-    bits = (__low) & 0xf;
-    out.z = float(*(__nv_fp4_e2m1*)(&bits));
-    bits = (__low >> 4) & 0xf;
-    out.w = float(*(__nv_fp4_e2m1*)(&bits));
-    return out;
-  }
-  uint8_t __high{0};
-  uint8_t __low{0};
-};

mlx/backend/cuda/quantized/fp_quantize.cu

@@ -1,18 +1,18 @@
 // 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/mxfp8_quantize.cuh"
 #include "mlx/backend/cuda/quantized/nvfp4_quantize.cuh"
 #include "mlx/backend/cuda/quantized/quantized.h"
-#include "mlx/backend/cuda/quantized/quantized_utils.cuh"
 #include "mlx/backend/cuda/vector_types.cuh"
 #include "mlx/dtype_utils.h"
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
-#include <cuda_fp4.h>
-#include <cuda_fp8.h>
+#include <cutlass/float8.h>
+#include <cutlass/numeric_conversion.h>
 
 constexpr float F8E4M3_MAX = 448.0f;
 constexpr float F4E2M1_MAX = 6.0f;
@@ -24,13 +24,29 @@ template <int bits>
 struct Dequantize {
   __device__ float operator()(uint8_t x) {
     if constexpr (bits == 8) {
-      return float(*(__nv_fp8_e4m3*)(&x));
+      return float(*(cutlass::float_e4m3_t*)(&x));
     } else {
-      return float(*(__nv_fp4_e2m1*)(&x));
+      return float(*(cutlass::float_e2m1_t*)(&x));
     }
   }
 };
 
+template <typename T>
+__device__ __forceinline__ void absmax_x2(T& out, const T& x1, const T& x2) {
+  if constexpr (
+      (std::is_same<T, __nv_bfloat162>::value) ||
+      (std::is_same<T, __half2>::value)) {
+    T a = x1;
+    T b = x2;
+    out = __hmax2(__habs2(a), __habs2(b));
+  } else if constexpr (std::is_same<T, float2>::value) {
+    float2 a = x1;
+    float2 b = x2;
+    out.x = fmaxf(fabsf(a.x), fabsf(b.x));
+    out.y = fmaxf(fabsf(a.y), fabsf(b.y));
+  }
+}
+
 namespace cg = cooperative_groups;
 
 template <typename T, int group_size, int bits, bool use_mx_scale, bool USE_SR>
@@ -45,7 +61,6 @@ __global__ void fp_quantize_dequantize(
   const float inv_scale_enc = use_global_scale ? 1.0f / scale_enc : 1.0f;
 
   using Tx2 = Vector2_t<T>;
-  using Tx4 = Vector4_t<T>;
   uint32_t rbits = 0; // reserved bits for future use
   auto block_size = cg::this_thread_block().dim_threads();
   auto block_idx = cg::this_thread_block().group_index();
@@ -79,30 +94,35 @@ __global__ void fp_quantize_dequantize(
   scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX;
   scale_dec_b *= scale_enc;
   // Convert to mx scale or nv scale
-  using ScaleType =
-      std::conditional_t<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
+  using ScaleType = std::conditional_t<
+      use_mx_scale,
+      cutlass::float_ue8m0_t,
+      cutlass::float_e4m3_t>;
   auto s = ScaleType(scale_dec_b);
   float scale_enc_b = scale_enc / float(s);
   float scale_dec = float(s) * inv_scale_enc;
   AlignedVector<T, group_size> w_hat;
 
 #pragma unroll
-  for (int i = 0; i < group_size / 4; i++) {
-    Tx4 w_Tx4 = *reinterpret_cast<Tx4*>(&w_tile[i * 4]);
-    float4 dq;
+  for (int i = 0; i < group_size / 8; i++) {
+    auto& w = *reinterpret_cast<cutlass::Array<T, 8>*>(&w_tile[i * 8]);
+    cutlass::NumericArrayConverter<float, T, 8> fp32_t;
+    auto scaled = fp32_t(w) * scale_enc_b;
+    cutlass::Array<float, 8> dq;
     if constexpr (bits == 8) {
-      uint32_t quantized_val =
-          scale_cvt_Tx4_to_fp8x4<T, USE_SR>(w_Tx4, scale_enc_b, rbits);
-      dq = dequant_fp8(quantized_val);
+      cutlass::NumericArrayConverter<cutlass::float_e4m3_t, float, 8> fp8_fp32;
+      auto quant = fp8_fp32(scaled);
+      cutlass::NumericArrayConverter<float, cutlass::float_e4m3_t, 8> fp32_fp8;
+      dq = fp32_fp8(quant);
     } else {
-      uint16_t quantized_val =
-          scale_cvt_Tx4_to_fp4x4<T, USE_SR>(w_Tx4, scale_enc_b, rbits);
-      dq = dequant_fp4(quantized_val);
+      cutlass::NumericArrayConverter<cutlass::float_e2m1_t, float, 8> fp4_fp32;
+      auto quant = fp4_fp32(scaled);
+      cutlass::NumericArrayConverter<float, cutlass::float_e2m1_t, 8> fp32_fp4;
+      dq = fp32_fp4(quant);
     }
-    w_hat[i * 4] = static_cast<T>(dq.x * scale_dec);
-    w_hat[i * 4 + 1] = static_cast<T>(dq.y * scale_dec);
-    w_hat[i * 4 + 2] = static_cast<T>(dq.z * scale_dec);
-    w_hat[i * 4 + 3] = static_cast<T>(dq.w * scale_dec);
+    cutlass::NumericArrayConverter<T, float, 8> t_fp32;
+    *reinterpret_cast<cutlass::Array<T, 8>*>(&w_hat[i * 8]) =
+        t_fp32(dq * scale_dec);
   }
   store_vector<group_size>(out, thread_idx, w_hat);
 }
@@ -157,10 +177,12 @@ __global__ void fp_quantize_rowwise(
   scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX;
   scale_dec_b *= scale_enc;
   // Convert to mx scale or nv scale
-  using ScaleType =
-      std::conditional_t<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
+  using ScaleType = std::conditional_t<
+      use_mx_scale,
+      cutlass::float_ue8m0_t,
+      cutlass::float_e4m3_t>;
   auto s = ScaleType(scale_dec_b);
-  uint8_t q_scale = s.__x;
+  uint8_t q_scale = s.storage;
   float scale_enc_b = scale_enc / float(s);
 
   scales[thread_idx] = q_scale;
@@ -256,11 +278,13 @@ __global__ void fp_quantize_columnwise(
     scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX;
     scale_dec_b *= scale_enc;
     // Convert to mx scale or nv scale
-    using ScaleType =
-        std::conditional_t<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
+    using ScaleType = std::conditional_t<
+        use_mx_scale,
+        cutlass::float_ue8m0_t,
+        cutlass::float_e4m3_t>;
     auto s = ScaleType(scale_dec_b);
     float scale_enc_b = scale_enc / float(s);
-    scales_smem[tidx][tidy] = s.__x;
+    scales_smem[tidx][tidy] = s.storage;
 
     int shared_idx = tidx * padded_local_cols + tidy * bytes_per_group;
 
@@ -345,8 +369,10 @@ __global__ void fp_dequantize(
   }
 
   size_t gindex = oindex / group_size;
-  using ScaleType =
-      std::conditional_t<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
+  using ScaleType = std::conditional_t<
+      use_mx_scale,
+      cutlass::float_ue8m0_t,
+      cutlass::float_e4m3_t>;
   auto scale = float(((ScaleType*)(scales))[gindex]) * inv_scale_enc;
 
   out += oindex;

mlx/backend/cuda/quantized/mxfp8_quantize.cuh

@@ -1,32 +1,23 @@
 #pragma once
 
-#include <cuda.h>
-#include <cuda_fp8.h>
-#include <cuda_runtime.h>
 #include "mlx/backend/cuda/vector_types.cuh"
 
-namespace mlx::core::cu {
+#include <cutlass/numeric_conversion.h>
 
-// TODO implement fast path
-template <typename T>
-__device__ __forceinline__ uint32_t
-scale_cvt_Tx4_to_fp8x4_fallback(const Vector4_t<T> input, const float scale) {
-  uint32_t out_fp8x4 = 0;
-  float4 scaled;
-  scaled.x = static_cast<float>(input.x) * scale;
-  scaled.y = static_cast<float>(input.y) * scale;
-  scaled.z = static_cast<float>(input.z) * scale;
-  scaled.w = static_cast<float>(input.w) * scale;
-  out_fp8x4 = __nv_fp8x4_e4m3(scaled).__x;
-  return out_fp8x4;
-}
+namespace mlx::core::cu {
 
 // Place holder for future fast path implementation
 template <typename T, bool USE_SR>
 __device__ __forceinline__ uint32_t scale_cvt_Tx4_to_fp8x4(
-    const Vector4_t<T> input,
+    const Vector4_t<T>& input,
     const float scale,
     uint32_t rbits) {
-  return scale_cvt_Tx4_to_fp8x4_fallback(input, scale);
+  cutlass::NumericArrayConverter<float, T, 4> fp32_t;
+  auto scaled =
+      fp32_t(*reinterpret_cast<const cutlass::Array<T, 4>*>(&input)) * scale;
+  cutlass::NumericArrayConverter<cutlass::float_e4m3_t, float, 4> fp8_fp32;
+  auto quant = fp8_fp32(scaled);
+  return *reinterpret_cast<uint32_t*>(&quant);
 }
-} // namespace mlx::core::cu
+
+} // namespace mlx::core::cu

mlx/backend/cuda/quantized/nvfp4_quantize.cuh

@@ -1,10 +1,9 @@
 #pragma once
 
-#include <cuda.h>
-#include <cuda_fp4.h>
-#include <cuda_runtime.h>
 #include "mlx/backend/cuda/vector_types.cuh"
 
+#include <cutlass/numeric_conversion.h>
+
 namespace mlx::core::cu {
 
 using bf16x4 = Vector4_t<__nv_bfloat16>;
@@ -13,23 +12,15 @@ using f32x4 = Vector4_t<float>;
 
 template <typename T>
 __device__ __forceinline__ uint16_t
-scale_cvt_Tx4_to_fp4x4_fallback(const Vector4_t<T> input, const float scale) {
+scale_cvt_Tx4_to_fp4x4_fallback(const Vector4_t<T>& input, const float scale) {
   // Fallback implementation for architectures that do not support cvt
   // instructions or for cuda versions with no fp4 support (< 12.8) -> scalar
-  uint16_t out_fp4x4 = 0;
-  fp32x4 scaled;
-  scaled.x = static_cast<float>(input.x) * scale;
-  scaled.y = static_cast<float>(input.y) * scale;
-  scaled.z = static_cast<float>(input.z) * scale;
-  scaled.w = static_cast<float>(input.w) * scale;
-  uint8_t q0 = __nv_fp4_e2m1(scaled.x).__x;
-  uint8_t q1 = __nv_fp4_e2m1(scaled.y).__x;
-  uint8_t q2 = __nv_fp4_e2m1(scaled.z).__x;
-  uint8_t q3 = __nv_fp4_e2m1(scaled.w).__x;
-  out_fp4x4 = (static_cast<uint16_t>(q3) << 12) |
-      (static_cast<uint16_t>(q2) << 8) | (static_cast<uint16_t>(q1) << 4) |
-      static_cast<uint16_t>(q0);
-  return out_fp4x4;
+  cutlass::NumericArrayConverter<float, T, 4> fp32_t;
+  auto scaled =
+      fp32_t(*reinterpret_cast<const cutlass::Array<T, 4>*>(&input)) * scale;
+  cutlass::NumericArrayConverter<cutlass::float_e2m1_t, float, 4> fp4_fp32;
+  auto quant = fp4_fp32(scaled);
+  return *reinterpret_cast<uint16_t*>(&quant);
 }
 
 #if (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) && \
@@ -318,7 +309,7 @@ __device__ __forceinline__ uint16_t scale_cvt_Tx4_to_fp4x4_fast(
 
 template <typename T, bool USE_SR>
 __device__ __forceinline__ uint16_t scale_cvt_Tx4_to_fp4x4(
-    const Vector4_t<T> input,
+    const Vector4_t<T>& input,
     const float scale,
     uint32_t rbits) {
 #if (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) && \
@@ -331,4 +322,5 @@ __device__ __forceinline__ uint16_t scale_cvt_Tx4_to_fp4x4(
   return scale_cvt_Tx4_to_fp4x4_fallback(input, scale);
 #endif
 }
-} // namespace mlx::core::cu
+
+} // namespace mlx::core::cu

mlx/backend/cuda/quantized/qmm/CMakeLists.txt

@@ -1,8 +1,35 @@
 target_sources(
   mlx
-  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/qmm.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n16_m1.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n32_m1.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n64_m2.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n128_m2.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n256_m2.cu)
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/qmm.cu ${CMAKE_CURRENT_SOURCE_DIR}/qmv.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/fp_qmv.cu)
+
+foreach(TileN 16 32 64 128 256)
+  set(OUTPUT_FILE "qmm_sm90_impl_n${TileN}.cu")
+  configure_file("${CMAKE_CURRENT_SOURCE_DIR}/qmm_sm90.cu"
+                 "${CMAKE_CURRENT_BINARY_DIR}/${OUTPUT_FILE}" @ONLY)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/${OUTPUT_FILE})
+endforeach()
+
+foreach(TileM 16 32 64)
+  set(OUTPUT_FILE "qmm_sm80_impl_m${TileM}.cu")
+  configure_file("${CMAKE_CURRENT_SOURCE_DIR}/qmm_sm80.cu"
+                 "${CMAKE_CURRENT_BINARY_DIR}/${OUTPUT_FILE}" @ONLY)
+  target_sources(mlx PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/${OUTPUT_FILE})
+endforeach()
+
+foreach(TileM 16 32 64)
+  foreach(KMajor true false)
+    foreach(HasKResidue true false)
+      foreach(SM80 true false)
+        if(${KMajor} AND ${HasKResidue})
+          continue()
+        endif()
+        set(OUTPUT_FILE
+            "qmm_naive_impl_m${TileM}_${KMajor}_${HasKResidue}_${SM80}.cu")
+        configure_file("${CMAKE_CURRENT_SOURCE_DIR}/qmm_naive.cu"
+                       "${CMAKE_CURRENT_BINARY_DIR}/${OUTPUT_FILE}" @ONLY)
+        target_sources(mlx PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/${OUTPUT_FILE})
+      endforeach()
+    endforeach()
+  endforeach()
+endforeach()

mlx/backend/cuda/quantized/qmm/cute_dequant.cuh

@@ -0,0 +1,169 @@
+// Copyright © 2026 Apple Inc.
+
+#pragma once
+
+#include <cute/numeric/numeric_types.hpp>
+#include <cute/tensor.hpp>
+#include <cutlass/numeric_conversion.h>
+#include <cuda/std/array>
+
+namespace cutlass {
+
+using uint3b_t = integer_subbyte<3, false>;
+using uint5b_t = integer_subbyte<5, false>;
+
+template <typename T, int N, FloatRoundStyle Round>
+struct NumericArrayConverter<T, uint3b_t, N, Round> {
+  static_assert(N % 8 == 0);
+
+  using result_type = Array<T, N>;
+  using source_type = Array<uint3b_t, N>;
+
+  CUTLASS_HOST_DEVICE
+  static result_type convert(const source_type& source) {
+    result_type result;
+    auto* s_base = reinterpret_cast<const uint8_t*>(&source);
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 8; ++i) {
+      auto* s = s_base + i * 3;
+      result[i * 8] = T(s[0] & 0x07);
+      result[i * 8 + 1] = T((s[0] & 0x38) >> 3);
+      result[i * 8 + 2] = T((s[0] & 0xc0) >> 6) + T((s[1] & 0x01) << 2);
+      result[i * 8 + 3] = T((s[1] & 0x0e) >> 1);
+      result[i * 8 + 4] = T((s[1] & 0x70) >> 4);
+      result[i * 8 + 5] = T((s[1] & 0x80) >> 7) + T((s[2] & 0x03) << 1);
+      result[i * 8 + 6] = T((s[2] & 0x1c) >> 2);
+      result[i * 8 + 7] = T((s[2] & 0xe0) >> 5);
+    }
+    return result;
+  }
+
+  CUTLASS_HOST_DEVICE
+  result_type operator()(const source_type& s) const {
+    return convert(s);
+  }
+};
+
+template <typename T, int N, FloatRoundStyle Round>
+struct NumericArrayConverter<T, uint5b_t, N, Round> {
+  static_assert(N % 8 == 0);
+
+  using result_type = Array<T, N>;
+  using source_type = Array<uint5b_t, N>;
+
+  CUTLASS_HOST_DEVICE
+  static result_type convert(const source_type& source) {
+    result_type result;
+    auto* s_base = reinterpret_cast<const uint8_t*>(&source);
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 8; ++i) {
+      auto* s = s_base + i * 5;
+      result[i * 8] = T(s[0] & 0x1f);
+      result[i * 8 + 1] = T((s[0] & 0xe0) >> 5) + T((s[1] & 0x03) << 3);
+      result[i * 8 + 2] = T((s[1] & 0x7c) >> 2);
+      result[i * 8 + 3] = T((s[1] & 0x80) >> 7) + T((s[2] & 0x0f) << 1);
+      result[i * 8 + 4] = T((s[2] & 0xf0) >> 4) + T((s[3] & 0x01) << 4);
+      result[i * 8 + 5] = T((s[3] & 0x3e) >> 1);
+      result[i * 8 + 6] = T((s[3] & 0xc0) >> 6) + T((s[4] & 0x07) << 2);
+      result[i * 8 + 7] = T((s[4] & 0xf8) >> 3);
+    }
+    return result;
+  }
+
+  CUTLASS_HOST_DEVICE
+  result_type operator()(const source_type& s) const {
+    return convert(s);
+  }
+};
+
+template <typename T, int N, FloatRoundStyle Round>
+struct NumericArrayConverter<T, uint6b_t, N, Round> {
+  static_assert(N % 4 == 0);
+
+  using result_type = Array<T, N>;
+  using source_type = Array<uint6b_t, N>;
+
+  CUTLASS_HOST_DEVICE
+  static result_type convert(const source_type& source) {
+    result_type result;
+    auto* s_base = reinterpret_cast<const uint8_t*>(&source);
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 4; ++i) {
+      auto* s = s_base + i * 3;
+      result[i * 4] = T(s[0] & 0x3f);
+      result[i * 4 + 1] = T((s[0] >> 6) & 0x03) + T((s[1] & 0x0f) << 2);
+      result[i * 4 + 2] = T((s[1] >> 4) & 0x0f) + T((s[2] & 0x03) << 4);
+      result[i * 4 + 3] = T((s[2] >> 2) & 0x3f);
+    }
+    return result;
+  }
+
+  CUTLASS_HOST_DEVICE
+  result_type operator()(const source_type& s) const {
+    return convert(s);
+  }
+};
+
+} // namespace cutlass
+
+namespace cute {
+
+// Required by tiled copy for 3/5/6-bit weights.
+struct uint24_t {
+  cuda::std::array<std::uint8_t, 3> bytes;
+};
+struct uint40_t {
+  cuda::std::array<std::uint8_t, 5> bytes;
+};
+struct uint48_t {
+  cuda::std::array<std::uint8_t, 6> bytes;
+};
+
+template <>
+struct uint_bit<24> {
+  using type = uint24_t;
+};
+template <>
+struct uint_bit<40> {
+  using type = uint40_t;
+};
+template <>
+struct uint_bit<48> {
+  using type = uint48_t;
+};
+
+} // namespace cute
+
+namespace cutlass_gemm {
+
+// Whether the quant type is affine quantization.
+template <typename Quant>
+constexpr bool quant_has_bias_v = !cutlass::has_negative_zero_v<Quant>;
+
+// Dequantize CuTe tensors with out = w * s + z.
+__device__ __forceinline__ void
+cute_vectorized_dequant(auto w, auto s, auto z, auto out) {
+  using namespace cute;
+  using Element = typename decltype(out)::value_type;
+  using Quant = typename decltype(w)::value_type;
+  // Scale must be one element.
+  CUTE_STATIC_ASSERT_V(cosize(s.layout()) == Int<1>{});
+  CUTE_STATIC_ASSERT_V(cosize(z.layout()) == Int<1>{});
+  // Quant must be contiguous.
+  auto layout = coalesce(w.layout());
+  CUTE_STATIC_ASSERT_V(stride(layout) == Int<1>{});
+  // Use cutlass for conversions.
+  constexpr int N = size(layout);
+  auto& w_vec = *(reinterpret_cast<const cutlass::Array<Quant, N>*>(
+      raw_pointer_cast(w.data())));
+  Element scale{s[0]};
+  cutlass::NumericArrayConverter<Element, Quant, N> converter;
+  auto w_dq = converter(w_vec) * scale;
+  if constexpr (quant_has_bias_v<Quant>) {
+    Element zero_point{z[0]};
+    w_dq = w_dq + zero_point;
+  }
+  copy(make_tensor(make_rmem_ptr<Element>(&w_dq), out.layout()), out);
+}
+
+} // namespace cutlass_gemm

mlx/backend/cuda/quantized/qmm/fp_qmv.cu

@@ -1,21 +1,25 @@
 // Copyright © 2025 Apple Inc.
 
+#include "mlx/backend/common/quantized.h"
 #include "mlx/backend/cuda/device/utils.cuh"
 #include "mlx/backend/cuda/kernel_utils.cuh"
-#include "mlx/backend/cuda/quantized/qmv.h"
-#include "mlx/backend/cuda/quantized/quantized_utils.cuh"
+#include "mlx/backend/cuda/quantized/qmm/qmm.h"
 #include "mlx/backend/cuda/quantized/quantized_utils.h"
 #include "mlx/dtype_utils.h"
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
+#include <cutlass/float8.h>
+#include <cutlass/numeric_conversion.h>
 
-namespace mlx::core::cu {
+namespace mlx::core {
+
+constexpr int rows_per_block = 8;
+
+namespace cu {
 
 namespace cg = cooperative_groups;
 
-static constexpr int rows_per_block = 8;
-
 template <typename T>
 __device__ void adjust_matrix_offsets(
     const T*& x,
@@ -74,8 +78,10 @@ __device__ void fp_qmv_impl(
   vec += g_idx.x * cols;
   out += g_idx.x * rows;
 
-  using ScaleType =
-      std::conditional_t<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
+  using ScaleType = std::conditional_t<
+      use_mx_scale,
+      cutlass::float_ue8m0_t,
+      cutlass::float_e4m3_t>;
   auto scales = (ScaleType*)(scales_);
   auto packed_cols = cols / vals_per_item;
 
@@ -100,35 +106,41 @@ __device__ void fp_qmv_impl(
         for (int j = 0; j < n_per_step; ++j) {
           int k = n_per_step * i + j;
           if constexpr (bits == 8) {
-            auto v = dequant_fp8(local_mat[k]);
+            cutlass::NumericArrayConverter<float, cutlass::float_e4m3_t, 4>
+                converter;
+            auto v = converter(
+                *reinterpret_cast<cutlass::Array<cutlass::float_e4m3_t, 4>*>(
+                    &local_mat[k]));
             local_sum.x +=
-                v.x * static_cast<float>(local_vec[vals_per_item * k]);
+                v[0] * static_cast<float>(local_vec[vals_per_item * k]);
             local_sum.x +=
-                v.y * static_cast<float>(local_vec[vals_per_item * k + 1]);
+                v[1] * static_cast<float>(local_vec[vals_per_item * k + 1]);
             local_sum.y +=
-                v.z * static_cast<float>(local_vec[vals_per_item * k + 2]);
+                v[2] * static_cast<float>(local_vec[vals_per_item * k + 2]);
             local_sum.y +=
-                v.w * static_cast<float>(local_vec[vals_per_item * k + 3]);
+                v[3] * static_cast<float>(local_vec[vals_per_item * k + 3]);
           } else {
-            auto v = dequant_fp4(local_mat[k]);
+            cutlass::NumericArrayConverter<float, cutlass::float_e2m1_t, 8>
+                converter;
+            auto v = converter(
+                *reinterpret_cast<cutlass::Array<cutlass::float_e2m1_t, 8>*>(
+                    &local_mat[k]));
             local_sum.x +=
-                v.x * static_cast<float>(local_vec[vals_per_item * k]);
+                v[0] * static_cast<float>(local_vec[vals_per_item * k]);
             local_sum.y +=
-                v.y * static_cast<float>(local_vec[vals_per_item * k + 1]);
+                v[1] * static_cast<float>(local_vec[vals_per_item * k + 1]);
             local_sum.x +=
-                v.z * static_cast<float>(local_vec[vals_per_item * k + 2]);
+                v[2] * static_cast<float>(local_vec[vals_per_item * k + 2]);
             local_sum.y +=
-                v.w * static_cast<float>(local_vec[vals_per_item * k + 3]);
-
-            v = dequant_fp4(local_mat[k] >> 16);
+                v[3] * static_cast<float>(local_vec[vals_per_item * k + 3]);
             local_sum.x +=
-                v.x * static_cast<float>(local_vec[vals_per_item * k + 4]);
+                v[4] * static_cast<float>(local_vec[vals_per_item * k + 4]);
             local_sum.y +=
-                v.y * static_cast<float>(local_vec[vals_per_item * k + 5]);
+                v[5] * static_cast<float>(local_vec[vals_per_item * k + 5]);
             local_sum.x +=
-                v.z * static_cast<float>(local_vec[vals_per_item * k + 6]);
+                v[6] * static_cast<float>(local_vec[vals_per_item * k + 6]);
             local_sum.y +=
-                v.w * static_cast<float>(local_vec[vals_per_item * k + 7]);
+                v[7] * static_cast<float>(local_vec[vals_per_item * k + 7]);
           }
         }
         sum += (local_sum.x + local_sum.y) * float(scales[i]);
@@ -199,6 +211,8 @@ __global__ void fp_qmv_batched(
       mat, scales, vec, out, rows, cols);
 }
 
+} // namespace cu
+
 template <typename F>
 void dispatch_1_2_4(int n, F&& f) {
   switch (n) {
@@ -221,11 +235,13 @@ void fp_qmv(
     array& out,
     int bits,
     int group_size,
-    int M,
-    int N,
-    int K,
-    CommandEncoder& encoder,
+    cu::CommandEncoder& encoder,
     Stream s) {
+  uint32_t M = x.shape(-2);
+  uint32_t N = out.shape(-1);
+  uint32_t K = x.shape(-1);
+  uint32_t B = out.size() / (M * N);
+
   // Make sure the last two dims of x and w, s, b are contiguous. This should
   // be relaxed for x.
   array vec = ensure_row_contiguous_matrix(x, encoder, s);
@@ -240,7 +256,6 @@ void fp_qmv(
     using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
     if constexpr (!std::is_same_v<T, double>) {
       dim3 block_dims{WARP_SIZE, rows_per_block};
-      uint32_t B = out.size() / (M * N);
       uint32_t blocks_y = (N + rows_per_block - 1) / rows_per_block;
       const uint32_t* mat_ptr = gpu_ptr<uint32_t>(mat);
       const T* vec_ptr = gpu_ptr<T>(vec);
@@ -256,55 +271,56 @@ void fp_qmv(
         n = 2;
       }
       dispatch_1_2_4(n, [&](auto n) {
-        dispatch_bool(B > 1, [&](auto batched) {
-          if (!batched.value) {
-            auto kernel =
-                fp_qmv_single<T, rows_per_block, n.value, 4, 32, true>;
-            if (bits == 8) {
-              kernel = fp_qmv_single<T, rows_per_block, n.value, 8, 32, true>;
-            } else if (group_size == 16) {
-              kernel = fp_qmv_single<T, rows_per_block, n.value, 4, 16, false>;
-            }
-            encoder.add_kernel_node(
-                kernel,
-                {static_cast<uint32_t>(M), blocks_y},
-                block_dims,
-                mat_ptr,
-                gpu_ptr<uint8_t>(scales),
-                vec_ptr,
-                gpu_ptr<T>(out),
-                N,
-                K);
-          } else {
-            auto kernel =
-                fp_qmv_batched<T, rows_per_block, n.value, 4, 32, true>;
-            if (bits == 8) {
-              kernel = fp_qmv_batched<T, rows_per_block, n.value, 8, 32, true>;
-            } else if (group_size == 16) {
-              kernel = fp_qmv_batched<T, rows_per_block, n.value, 4, 16, false>;
-            }
-            encoder.add_kernel_node(
-                kernel,
-                {static_cast<uint32_t>(M), blocks_y, B},
-                block_dims,
-                mat_ptr,
-                gpu_ptr<uint8_t>(scales),
-                vec_ptr,
-                gpu_ptr<T>(out),
-                N,
-                K,
-                vec.ndim() - 2,
-                const_param(vec.shape()),
-                const_param(vec.strides()),
-                mat.ndim() - 2,
-                const_param(mat.shape()),
-                const_param(mat.strides()),
-                const_param(scales.strides()));
+        if (B == 1) {
+          auto kernel =
+              cu::fp_qmv_single<T, rows_per_block, n.value, 4, 32, true>;
+          if (bits == 8) {
+            kernel = cu::fp_qmv_single<T, rows_per_block, n.value, 8, 32, true>;
+          } else if (group_size == 16) {
+            kernel =
+                cu::fp_qmv_single<T, rows_per_block, n.value, 4, 16, false>;
           }
-        });
+          encoder.add_kernel_node(
+              kernel,
+              {uint32_t(x.size() / K), blocks_y},
+              block_dims,
+              mat_ptr,
+              gpu_ptr<uint8_t>(scales),
+              vec_ptr,
+              gpu_ptr<T>(out),
+              N,
+              K);
+        } else {
+          auto kernel =
+              cu::fp_qmv_batched<T, rows_per_block, n.value, 4, 32, true>;
+          if (bits == 8) {
+            kernel =
+                cu::fp_qmv_batched<T, rows_per_block, n.value, 8, 32, true>;
+          } else if (group_size == 16) {
+            kernel =
+                cu::fp_qmv_batched<T, rows_per_block, n.value, 4, 16, false>;
+          }
+          encoder.add_kernel_node(
+              kernel,
+              {M, blocks_y, B},
+              block_dims,
+              mat_ptr,
+              gpu_ptr<uint8_t>(scales),
+              vec_ptr,
+              gpu_ptr<T>(out),
+              N,
+              K,
+              vec.ndim() - 2,
+              const_param(vec.shape()),
+              const_param(vec.strides()),
+              mat.ndim() - 2,
+              const_param(mat.shape()),
+              const_param(mat.strides()),
+              const_param(scales.strides()));
+        }
       });
     }
   });
 }
 
-} // namespace mlx::core::cu
+} // namespace mlx::core

mlx/backend/cuda/quantized/qmm/qmm.cpp

@@ -1,60 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm.h"
-
-#include <cute/tensor.hpp>
-
-namespace mlx::core {
-
-#if defined(MLX_CUDA_SM90A_ENABLED)
-// Defined in qmm_impl_sm90_xxx.cu files.
-template <typename TileShape, typename ClusterShape>
-void qmm_impl_sm90(
-    const array& x,
-    const array& w,
-    const array& scales,
-    const array& biases,
-    array& out,
-    int bits,
-    int group_size,
-    cu::CommandEncoder& encoder,
-    Stream s);
-#endif // defined(MLX_CUDA_SM90A_ENABLED)
-
-void qmm_sm90(
-    const array& x,
-    const array& w,
-    const array& scales,
-    const array& biases,
-    array& out,
-    int bits,
-    int group_size,
-    cu::CommandEncoder& encoder,
-    Stream s) {
-#if defined(MLX_CUDA_SM90A_ENABLED)
-  auto dispatch = [&]<int tile_m, int tile_n, int cluster_m>() {
-    using cute::Int;
-    using TileShapeMN = cute::Shape<Int<tile_m>, Int<tile_n>>;
-    using ClusterShape = cute::Shape<Int<cluster_m>, Int<1>, Int<1>>;
-    qmm_impl_sm90<TileShapeMN, ClusterShape>(
-        x, w, scales, biases, out, bits, group_size, encoder, s);
-  };
-  int m = out.shape(-2);
-  if (m <= 16) {
-    dispatch.template operator()<128, 16, 1>();
-  } else if (m <= 32) {
-    dispatch.template operator()<128, 32, 1>();
-  } else if (m <= 64) {
-    dispatch.template operator()<128, 64, 2>();
-  } else if (m <= 128) {
-    dispatch.template operator()<128, 128, 2>();
-  } else {
-    dispatch.template operator()<128, 256, 2>();
-  }
-#else
-  throw std::runtime_error(
-      "[quantized_matmul] Hopper-only kernel is not available.");
-#endif // defined(MLX_CUDA_SM90A_ENABLED)
-}
-
-} // namespace mlx::core

mlx/backend/cuda/quantized/qmm/qmm.cu

@@ -0,0 +1,359 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/cuda/quantized/qmm/qmm.h"
+
+#include <cute/tensor.hpp>
+
+namespace mlx::core {
+
+namespace {
+
+inline bool is_last_2_dims_row_contiguous(const array& x) {
+  return x.flags().contiguous && (x.ndim() >= 2) && (x.strides(-1) == 1) &&
+      (x.strides(-2) == x.shape(-1));
+}
+
+} // namespace
+
+#if defined(MLX_CUDA_SM90A_ENABLED)
+// Defined in qmm_sm90.cu.
+template <int TileN>
+void qmm_sm90_impl(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const array& biases,
+    array& out,
+    int bits,
+    int group_size,
+    cu::CommandEncoder& encoder,
+    Stream s);
+#endif // defined(MLX_CUDA_SM90A_ENABLED)
+
+bool supports_qmm_sm90(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const array& out,
+    bool transpose,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::Device& device) {
+  if (device.compute_capability_major() != 9) {
+    return false;
+  }
+  int k = x.shape(-1);
+  if (k % 64 != 0) {
+    return false;
+  }
+  if (!biases) {
+    return false;
+  }
+  if (!is_last_2_dims_row_contiguous(w) ||
+      !is_last_2_dims_row_contiguous(scales) ||
+      !is_last_2_dims_row_contiguous(*biases)) {
+    return false;
+  }
+  if (!transpose) {
+    return false;
+  }
+  if (bits % 2 != 0) {
+    return false;
+  }
+  if (group_size < k) {
+    return false;
+  }
+  if (mode != QuantizationMode::Affine) {
+    return false;
+  }
+  return true;
+}
+
+void qmm_sm90(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const array& biases,
+    array& out,
+    int bits,
+    int group_size,
+    cu::CommandEncoder& encoder,
+    Stream s) {
+#if defined(MLX_CUDA_SM90A_ENABLED)
+  auto dispatch = [&]<int TileN>() {
+    qmm_sm90_impl<TileN>(
+        x, w, scales, biases, out, bits, group_size, encoder, s);
+  };
+  int m = out.ndim() > 1 ? out.shape(-2) : 1;
+  if (m <= 16) {
+    dispatch.template operator()<16>();
+  } else if (m <= 32) {
+    dispatch.template operator()<32>();
+  } else if (m <= 64) {
+    dispatch.template operator()<64>();
+  } else if (m <= 128) {
+    dispatch.template operator()<128>();
+  } else {
+    dispatch.template operator()<256>();
+  }
+#else
+  throw std::runtime_error(
+      "[quantized_matmul] Hopper-only kernel is not available.");
+#endif // defined(MLX_CUDA_SM90A_ENABLED)
+}
+
+// Defined in qmm_sm80.cu.
+template <int TileM>
+void qmm_sm80_impl(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const std::optional<array>& lhs_indices,
+    const std::optional<array>& rhs_indices,
+    array& out,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::CommandEncoder& encoder);
+
+bool supports_qmm_sm80(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const array& out,
+    bool transpose,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::Device& device) {
+  if (device.compute_capability_major() < 8) {
+    return false;
+  }
+  int n = out.shape(-1);
+  int k = x.shape(-1);
+  if ((n % 128 != 0) || (k % std::max(64, group_size) != 0)) {
+    return false;
+  }
+  if (!is_last_2_dims_row_contiguous(w) ||
+      !is_last_2_dims_row_contiguous(scales)) {
+    return false;
+  }
+  if (biases && !is_last_2_dims_row_contiguous(*biases)) {
+    return false;
+  }
+  if (x.dtype() != float16 && x.dtype() != bfloat16) {
+    return false;
+  }
+  if (!transpose) {
+    return false;
+  }
+  if (bits != 4 && bits != 8) {
+    return false;
+  }
+  return true;
+}
+
+void qmm_sm80(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const std::optional<array>& lhs_indices,
+    const std::optional<array>& rhs_indices,
+    array& out,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::CommandEncoder& encoder) {
+  auto dispatch = [&]<int TileM>() {
+    qmm_sm80_impl<TileM>(
+        x,
+        w,
+        scales,
+        biases,
+        lhs_indices,
+        rhs_indices,
+        out,
+        bits,
+        group_size,
+        mode,
+        encoder);
+  };
+  int m = out.ndim() > 1 ? out.shape(-2) : 1;
+  if (m <= 16) {
+    dispatch.template operator()<16>();
+  } else if (m <= 32) {
+    dispatch.template operator()<32>();
+  } else {
+    dispatch.template operator()<64>();
+  }
+}
+
+// Defined in qmm_naive.cu.
+template <int TileM, bool KMajor, bool HasKResidue, bool SM80>
+void qmm_naive_impl(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const std::optional<array>& lhs_indices,
+    const std::optional<array>& rhs_indices,
+    array& out,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::CommandEncoder& encoder);
+
+bool supports_qmm_naive(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const array& out,
+    bool transpose,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::Device& device) {
+  int k = x.shape(-1);
+  if (transpose && (k % std::max(64, group_size) != 0)) {
+    return false;
+  }
+  if (!is_last_2_dims_row_contiguous(w) ||
+      !is_last_2_dims_row_contiguous(scales)) {
+    return false;
+  }
+  if (biases && !is_last_2_dims_row_contiguous(*biases)) {
+    return false;
+  }
+  return true;
+}
+
+void qmm_naive(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const std::optional<array>& lhs_indices,
+    const std::optional<array>& rhs_indices,
+    array& out,
+    bool transpose,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::CommandEncoder& encoder) {
+  auto dispatch = [&]<int TileM, bool KMajor, bool HasKResidue, bool SM80>() {
+    qmm_naive_impl<TileM, KMajor, HasKResidue, SM80>(
+        x,
+        w,
+        scales,
+        biases,
+        lhs_indices,
+        rhs_indices,
+        out,
+        bits,
+        group_size,
+        mode,
+        encoder);
+  };
+  auto dispatch_k = [&](auto k_major, bool has_k_residue, auto&& f) {
+    if constexpr (k_major.value) {
+      if (has_k_residue) {
+        throw std::invalid_argument(
+            "[quantized_matmul] K must be multiples of max(64, group_size).");
+      }
+      f.template operator()<false>();
+    } else {
+      dispatch_bool(has_k_residue, [&](auto has_k_residue) {
+        f.template operator()<has_k_residue.value>();
+      });
+    }
+  };
+  int m = out.ndim() > 1 ? out.shape(-2) : 1;
+  int k = x.shape(-1);
+  int tile_k = std::max(64, group_size);
+  bool has_k_residue = k % tile_k != 0;
+  bool sm80 = encoder.device().compute_capability_major() >= 8;
+  dispatch_bool(transpose, [&](auto k_major) {
+    dispatch_k(k_major, has_k_residue, [&]<bool HasKResidue>() {
+      dispatch_bool(sm80, [&](auto sm80) {
+        constexpr bool KMajor = k_major.value;
+        constexpr bool SM80 = sm80.value;
+        if (m <= 16) {
+          dispatch.template operator()<16, KMajor, HasKResidue, SM80>();
+        } else if (m <= 32) {
+          dispatch.template operator()<32, KMajor, HasKResidue, SM80>();
+        } else {
+          dispatch.template operator()<64, KMajor, HasKResidue, SM80>();
+        }
+      });
+    });
+  });
+}
+
+bool supports_fp_qmv(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const array& out,
+    bool transpose,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::Device& device) {
+  // The fp_qmv kernel uses less registers and is faster for sm120. For sm80/90
+  // the qmv kernel is faster. We didn't test sm89/100.
+  if (device.compute_capability_major() <= 9) {
+    return false;
+  }
+  bool non_batched = w.ndim() == 2;
+  int k = x.shape(-1);
+  int n = out.shape(-1);
+  int vec_batch = non_batched ? x.size() / k : x.shape(-2);
+  if (vec_batch > 8) {
+    return false;
+  }
+  if (!transpose) {
+    return false;
+  }
+  if (mode == QuantizationMode::Affine) {
+    return false;
+  }
+  return true;
+}
+
+bool supports_qmv(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const array& out,
+    bool transpose,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::Device& device) {
+  int k = x.shape(-1);
+  if (k % 8 != 0) {
+    return false;
+  }
+  if (!is_last_2_dims_row_contiguous(w) ||
+      !is_last_2_dims_row_contiguous(scales)) {
+    return false;
+  }
+  if (biases && !is_last_2_dims_row_contiguous(*biases)) {
+    return false;
+  }
+  if (!transpose) {
+    return false;
+  }
+  return true;
+}
+
+} // namespace mlx::core

mlx/backend/cuda/quantized/qmm/qmm.h

@@ -3,11 +3,24 @@
 #pragma once
 
 #include "mlx/backend/cuda/device.h"
+#include "mlx/primitives.h"
 
 #include <optional>
 
 namespace mlx::core {
 
+bool supports_qmm_sm90(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const array& out,
+    bool transpose,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::Device& device);
+
 void qmm_sm90(
     const array& x,
     const array& w,
@@ -19,4 +32,113 @@ void qmm_sm90(
     cu::CommandEncoder& encoder,
     Stream s);
 
+bool supports_qmm_sm80(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const array& out,
+    bool transpose,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::Device& device);
+
+void qmm_sm80(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const std::optional<array>& lhs_indices,
+    const std::optional<array>& rhs_indices,
+    array& out,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::CommandEncoder& encoder);
+
+bool supports_qmm_naive(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const array& out,
+    bool transpose,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::Device& device);
+
+void qmm_naive(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const std::optional<array>& lhs_indices,
+    const std::optional<array>& rhs_indices,
+    array& out,
+    bool transpose,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::CommandEncoder& encoder);
+
+bool supports_fp_qmv(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const array& out,
+    bool transpose,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::Device& device);
+
+void fp_qmv(
+    const array& x,
+    const array& w,
+    const array& scales,
+    array& out,
+    int bits,
+    int group_size,
+    cu::CommandEncoder& encoder,
+    Stream s);
+
+bool supports_qmv(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const array& out,
+    bool transpose,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::Device& device);
+
+void qmv(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    array& out,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::CommandEncoder& encoder);
+
+void gather_qmv(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const array& lhs_indices,
+    const array& rhs_indices,
+    array& out,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::CommandEncoder& encoder);
+
 } // namespace mlx::core

mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n128_m2.cu

@@ -1,10 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh"
-
-using namespace cute;
-
-using TileShapeMN = Shape<_128, _128>;
-using ClusterShape = Shape<_2, _1, _1>;
-
-QMM_SM90_GPU(TileShapeMN, ClusterShape)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n16_m1.cu

@@ -1,10 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh"
-
-using namespace cute;
-
-using TileShapeMN = Shape<_128, _16>;
-using ClusterShape = Shape<_1, _1, _1>;
-
-QMM_SM90_GPU(TileShapeMN, ClusterShape)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n256_m2.cu

@@ -1,10 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh"
-
-using namespace cute;
-
-using TileShapeMN = Shape<_128, _256>;
-using ClusterShape = Shape<_2, _1, _1>;
-
-QMM_SM90_GPU(TileShapeMN, ClusterShape)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n32_m1.cu

@@ -1,10 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh"
-
-using namespace cute;
-
-using TileShapeMN = Shape<_128, _32>;
-using ClusterShape = Shape<_1, _1, _1>;
-
-QMM_SM90_GPU(TileShapeMN, ClusterShape)

mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n64_m2.cu

@@ -1,10 +0,0 @@
-// Copyright © 2026 Apple Inc.
-
-#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh"
-
-using namespace cute;
-
-using TileShapeMN = Shape<_128, _64>;
-using ClusterShape = Shape<_2, _1, _1>;
-
-QMM_SM90_GPU(TileShapeMN, ClusterShape)

mlx/backend/cuda/quantized/qmm/qmm_naive.cu

@@ -0,0 +1,244 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/cuda/quantized/qmm/qmm.h"
+#include "mlx/backend/cuda/quantized/qmm/qmm_naive.cuh"
+
+// clang-format off
+
+// We can't put kernel code in mlx::core due to name conflicts of "Shape".
+namespace cutlass_gemm {
+
+using namespace cute;
+
+template <bool KMajor, bool HasKResidue, bool SM80,
+          typename Element, typename Quant, typename Scale,
+          typename ProblemShape,
+          typename CtaTiler,
+          typename StrideA,
+          typename StrideB,
+          typename LayoutS,
+          typename StrideC,
+          typename TiledMma>
+__global__
+__launch_bounds__(decltype(size(TiledMma{}))::value)
+void qmm_naive_kernel(
+    ProblemShape shape_MNKL,
+    CtaTiler cta_tiler,
+    const Element* A, StrideA dA,
+    const Quant* B, StrideB dB,
+    const Scale* S, const Element* Z, LayoutS S_layout,
+    const uint32_t* lhs_indices, const uint32_t* rhs_indices,
+    Element* C, StrideC dC,
+    TiledMma mma) {
+  CUTE_STATIC_ASSERT_V(congruent(select<0,2,3>(shape_MNKL), dA));
+  CUTE_STATIC_ASSERT_V(congruent(select<1,2,3>(shape_MNKL), dB));
+  CUTE_STATIC_ASSERT_V(congruent(select<0,1,3>(shape_MNKL), dC));
+
+  int thread_idx = int(threadIdx.x);
+  int m_coord = int(blockIdx.x);
+  int n_coord = int(blockIdx.y);
+  int l_coord = int(blockIdx.z);
+
+  // Represent the full tensors.
+  Tensor mA_mkl = make_tensor(make_gmem_ptr(A),        select<0,2,3>(shape_MNKL), dA); // (M,K,L)
+  Tensor mB_nkl = make_tensor(make_gmem_ptr<Quant>(B), select<1,2,3>(shape_MNKL), dB); // (N,K,L)
+  Tensor mC_mnl = make_tensor(make_gmem_ptr(C),        select<0,1,3>(shape_MNKL), dC); // (M,N,L)
+
+  Tensor mS_nkl = make_tensor(make_gmem_ptr(S), S_layout); // (N,(group_size,K/group_size),L)
+  Tensor mZ_nkl = make_tensor(make_gmem_ptr(Z), S_layout); // (N,(group_size,K/group_size),L)
+
+  // For gather, use index lookup for input batch slicing.
+  uint32_t a_batch = lhs_indices ? lhs_indices[l_coord] : l_coord;
+  uint32_t b_batch = rhs_indices ? rhs_indices[l_coord] : l_coord;
+
+  // Get batch slice.
+  Tensor mA = mA_mkl(_,_,a_batch); // (M,K)
+  Tensor mB = mB_nkl(_,_,b_batch); // (N,K)
+  Tensor mC = mC_mnl(_,_,l_coord); // (M,N)
+
+  Tensor mS = mS_nkl(_,_,b_batch); // (N,(group_size,K/group_size))
+  Tensor mZ = mZ_nkl(_,_,b_batch); // (N,(group_size,K/group_size))
+
+  // Get the appropriate blocks for this thread block.
+  auto cta_coord = make_coord(m_coord, n_coord, _); // (m,n,k)
+  Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k)
+  Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
+  Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N)
+
+  Tensor gS = local_tile(mS, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
+  Tensor gZ = local_tile(mZ, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
+
+  // Compute tile residues for predication.
+  int m_max_coord = size<0>(shape_MNKL) - size<0>(cta_tiler) * m_coord; // M - BLK_M * m_coord
+  int n_max_coord = size<1>(shape_MNKL) - size<1>(cta_tiler) * n_coord; // N - BLK_N * n_coord
+  int k_residue = size<2>(shape_MNKL) - size<1>(gA) * size<2>(gA);
+
+  qmm_naive_mainloop<KMajor, HasKResidue, SM80>(
+      cta_tiler,
+      gA,
+      gB,
+      gS,
+      gZ,
+      gC,
+      mma,
+      m_max_coord, n_max_coord, k_residue,
+      thread_idx);
+}
+
+template <int TileM, bool KMajor, bool HasKResidue, bool SM80,
+          typename Element, typename Quant, typename Scale>
+void qmm_naive(
+    const Element* A,
+    const Quant*   B,
+    const Scale*   S,
+    const Element* Z,
+    const uint32_t* lhs_indices,
+    const uint32_t* rhs_indices,
+    Element* C,
+    int m, int n, int k, int l,
+    bool broadcast_b,
+    auto group_size,
+    auto&& launch_kernel) {
+  // Define shapes (dynamic).
+  auto shape_MNKL = make_shape(m, n, k, l); // (M,N,K,L)
+
+  // Define layouts (mixed).
+  auto dA = make_stride(k, Int<1>{}, m * k);  // (dM,dK,dL)
+  auto dB = make_matrix_stride<KMajor>(n, k); // (dN,dK,dL)
+  auto dC = make_stride(n, Int<1>{}, m * n);  // (dM,dN,dL)
+  auto S_layout = make_scales_layout<KMajor>(n, k, l, group_size);
+
+  // Handle broadcasting.
+  if (broadcast_b) {
+    get<2>(dB) = 0;
+    get<2>(stride(S_layout)) = 0;
+  }
+
+  // Define CTA tile size (static).
+  auto cta_tiler = make_cta_tiler<TileM, SM80>(group_size);
+
+  // Define MMA.
+  auto mma = make_tiled_mma<SM80, Element>(cta_tiler);
+  auto num_threads = size(mma);
+
+  // Shared memory size.
+  auto [sA_layout, sB_layout] = make_smem_layouts<KMajor>(cta_tiler);
+  size_t smem_bytes = sizeof(SharedStorage<Element, decltype(sA_layout), decltype(sB_layout)>);
+
+  auto* kernel = &qmm_naive_kernel<
+      KMajor, HasKResidue, SM80,
+      Element, Quant, Scale,
+      decltype(shape_MNKL),
+      decltype(cta_tiler),
+      decltype(dA),
+      decltype(dB),
+      decltype(S_layout),
+      decltype(dC),
+      decltype(mma)>;
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
+
+  dim3 num_blocks{uint32_t(ceil_div(m, size<0>(cta_tiler))),
+                  uint32_t(ceil_div(n, size<1>(cta_tiler))),
+                  uint32_t(l)};
+  dim3 block_dims{num_threads};
+  void* args[] = {
+      &shape_MNKL,
+      &cta_tiler,
+      &A, &dA,
+      &B, &dB,
+      &S, &Z, &S_layout,
+      &lhs_indices, &rhs_indices,
+      &C, &dC,
+      &mma};
+  launch_kernel(reinterpret_cast<void*>(kernel), num_blocks, block_dims, smem_bytes, args);
+}
+
+} // namespace cutlass_gemm
+
+// clang-format on
+
+namespace mlx::core {
+
+template <int TileM, bool KMajor, bool HasKResidue, bool SM80>
+void qmm_naive_impl(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const std::optional<array>& lhs_indices,
+    const std::optional<array>& rhs_indices,
+    array& out,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::CommandEncoder& encoder) {
+  const char* tag = "[quantized_matmul]";
+  int m = out.ndim() > 1 ? out.shape(-2) : 1;
+  int n = out.shape(-1);
+  int k = x.shape(-1);
+  int l = out.size() / (m * n);
+  bool broadcast_b = (w.ndim() <= 2) || (w.size() != w.data_size());
+
+  dispatch_element_types(out.dtype(), tag, [&]<typename Element>() {
+    dispatch_quant_types<Element>(
+        bits,
+        group_size,
+        mode,
+        tag,
+        [&]<typename Quant, typename Scale, int group_size>() {
+          encoder.set_input_array(x);
+          encoder.set_input_array(w);
+          encoder.set_input_array(scales);
+          if (biases) {
+            encoder.set_input_array(*biases);
+          }
+          if (lhs_indices) {
+            encoder.set_input_array(*lhs_indices);
+          }
+          if (rhs_indices) {
+            encoder.set_input_array(*rhs_indices);
+          }
+          encoder.set_output_array(out);
+          cutlass_gemm::qmm_naive<TileM, KMajor, HasKResidue, SM80>(
+              gpu_ptr<Element>(x),
+              gpu_ptr<Quant>(w),
+              gpu_ptr<Scale>(scales),
+              biases ? gpu_ptr<Element>(*biases) : nullptr,
+              lhs_indices ? gpu_ptr<uint32_t>(*lhs_indices) : nullptr,
+              rhs_indices ? gpu_ptr<uint32_t>(*rhs_indices) : nullptr,
+              gpu_ptr<Element>(out),
+              m,
+              n,
+              k,
+              l,
+              broadcast_b,
+              cute::Int<group_size>{},
+              [&](auto* kernel,
+                  dim3 num_blocks,
+                  dim3 block_dims,
+                  size_t smem_bytes,
+                  void** args) {
+                encoder.add_kernel_node_raw(
+                    kernel, num_blocks, block_dims, {}, smem_bytes, args);
+              });
+        });
+  });
+}
+
+// clang-format off
+template void qmm_naive_impl<@TileM@, @KMajor@, @HasKResidue@, @SM80@>(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const std::optional<array>& lhs_indices,
+    const std::optional<array>& rhs_indices,
+    array& out,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::CommandEncoder& encoder);
+// clang-format on
+
+} // namespace mlx::core

mlx/backend/cuda/quantized/qmm/qmm_naive.cuh

@@ -0,0 +1,381 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/quantized/qmm/cute_dequant.cuh"
+#include "mlx/dtype_utils.h"
+
+// clang-format off
+
+// We can't put kernel code in mlx::core due to name conflicts of "Shape".
+namespace cutlass_gemm {
+
+using namespace cute;
+
+template <typename Element, typename SmemLayoutA, typename SmemLayoutB>
+struct SharedStorage {
+  ArrayEngine<Element, cosize_v<SmemLayoutA>> A;
+  ArrayEngine<Element, cosize_v<SmemLayoutB>> B;
+};
+
+template <bool KMajor = true>
+inline constexpr auto make_smem_layout(auto bM, auto bK) {
+  // TODO: Calculate swizzle based on tile shape.
+  if constexpr (KMajor) {
+    auto swizzle = composition(Swizzle<3,3,3>{},
+                               Layout<Shape <_8,Shape <_8, _8>>,
+                                      Stride<_8,Stride<_1,_64>>>{});
+    return tile_to_shape(swizzle, make_shape(bM, bK));
+  } else {
+    auto swizzle = composition(Swizzle<3,3,3>{},
+                               Layout<Shape<_64,_1>, Stride<_1,_64>>{});
+    return tile_to_shape(swizzle, make_shape(bM, bK));
+  }
+}
+
+template <bool KMajor = true>
+inline constexpr auto make_smem_layouts(auto cta_tiler) {
+  auto [bM, bN, bK] = cta_tiler;
+  auto sA_layout = make_smem_layout(bM, bK);
+  auto sB_layout = make_smem_layout<KMajor>(bN, bK);
+  return std::make_tuple(sA_layout, sB_layout);
+}
+
+template <typename T, bool KMajor = true, bool HasKResidue = false>
+inline constexpr auto make_tiled_copy(auto num_threads, auto bM, auto bK) {
+  // TODO: Only do 1-element read for the tile of residue.
+  auto n_read = Int<HasKResidue ? 1 : 8>{};
+  auto atom = Copy_Atom<UniversalCopy<uint_bit_t<n_read * sizeof_bits_v<T>>>, T>{};
+  if constexpr (KMajor) {
+    auto k_threads = bK / n_read;
+    return make_tiled_copy(
+        atom,
+        make_layout(make_shape(Int<num_threads / k_threads>{}, k_threads), LayoutRight{}),
+        make_layout(make_shape(Int<1>{}, n_read)));
+  } else {
+    auto m_threads = bM / n_read;
+    return make_tiled_copy(
+        atom,
+        make_layout(make_shape(m_threads, Int<num_threads / m_threads>{}), LayoutLeft{}),
+        make_layout(make_shape(n_read, Int<1>{})));
+  }
+}
+
+
+__device__ __forceinline__ void
+cute_naive_dequant(auto w, auto s, auto z, auto out) {
+  using Element = typename decltype(out)::value_type;
+  using Quant = typename decltype(w)::value_type;
+  using Scale = typename decltype(s)::value_type;
+  transform(w, out, [](Quant q) { return Element(q); } );
+  transform(out, s, out, [](Element e, Scale s) { return e * Element(s); });
+  if constexpr (quant_has_bias_v<Quant>) {
+    transform(out, z, out, plus{});
+  }
+}
+
+__device__ __forceinline__ void
+cute_dequant(auto w, auto s, auto z, auto out) {
+  if constexpr (stride(coalesce(w.layout())) == Int<1>{} &&
+                is_static_v<decltype(s.layout())>) {
+    cute_vectorized_dequant(w, s, z, out);
+  } else {
+    cute_naive_dequant(w, s, z, out);
+  }
+}
+
+template <bool KMajor, bool HasKResidue, bool SM80,
+          typename CtaTiler,
+          typename TensorA,
+          typename TensorB,
+          typename TensorS,
+          typename TensorZ,
+          typename TensorC,
+          typename TiledMma>
+CUTE_DEVICE void qmm_naive_mainloop(
+    CtaTiler cta_tiler,
+    TensorA gA,
+    TensorB gB,
+    TensorS gS,
+    TensorZ gZ,
+    TensorC gC,
+    TiledMma mma,
+    int m_max_coord,
+    int n_max_coord,
+    int k_residue,
+    int thread_idx) {
+  // Get the types of operands.
+  using Element = decltype(gA)::value_type;
+  using Quant = decltype(gB)::value_type;
+
+  // Shift tensor so we handle residue of K in the 0th tile.
+  gA = domain_offset(make_coord(0, k_residue, 0), gA);
+  if constexpr (sizeof_bits_v<Quant> % 8 == 0) {
+    gB = domain_offset(make_coord(0, k_residue, 0), gB);
+  } else {
+    gB.data() = recast_ptr<Quant>(raw_pointer_cast(gB.data()) + gB.layout()(0, k_residue, 0) * cuda::std::min(8, sizeof_bits_v<Quant>) / 8);
+  }
+  gS = domain_offset(make_coord(0, k_residue, 0), gS);
+  gZ = domain_offset(make_coord(0, k_residue, 0), gZ);
+
+  // Define smem layouts.
+  auto [sA_layout, sB_layout] = make_smem_layouts(cta_tiler);
+
+  // Shared memory buffer.
+  extern __shared__ char smem_buf[];
+  using SharedStorage = SharedStorage<Element, decltype(sA_layout), decltype(sB_layout)>;
+  SharedStorage& smem = *reinterpret_cast<SharedStorage*>(smem_buf);
+  Tensor sA = make_tensor(make_smem_ptr(smem.A.begin()), sA_layout); // (BLK_M,BLK_K)
+  Tensor sB = make_tensor(make_smem_ptr(smem.B.begin()), sB_layout); // (BLK_N,BLK_K)
+
+  // Define copy atoms.
+  auto num_threads = size(mma);
+  auto [bM, bN, bK] = cta_tiler;
+  TiledCopy copy_a = make_tiled_copy<Element, true, HasKResidue>(num_threads, bM, bK);
+  TiledCopy copy_b = make_tiled_copy<Quant, KMajor>(num_threads, bN, bK);
+
+  // Partition the copying of A/B/C tiles across the threads.
+  ThrCopy thr_copy_a = copy_a.get_slice(thread_idx);
+  Tensor tAgA = thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k)
+  Tensor tAsA = thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K)
+  Tensor tArA = make_fragment_like(tAsA);   // (ACPY,ACPY_M,ACPY_K)
+
+  ThrCopy thr_copy_b = copy_b.get_slice(thread_idx);
+  Tensor tBgB = thr_copy_b.partition_S(gB);        // (BCPY,BCPY_N,BCPY_K,k)
+  Tensor tBsB = thr_copy_b.partition_D(sB);        // (BCPY,BCPY_N,BCPY_K)
+  Tensor tBrB = make_fragment_like<Quant>(tBsB);   // (BCPY,BCPY_M,BCPY_K)
+  Tensor tBrB_dq = make_fragment_like(tBsB);       // (BCPY,BCPY_M,BCPY_K)
+  Tensor tBgS = thr_copy_b.partition_S(gS);        // (BCPY,BCPY_N,BCPY_K,k)
+  Tensor tBrS = make_fragment_like(tBgS(_,_,_,0)); // (BCPY,BCPY_N,BCPY_K)
+  Tensor tBgZ = thr_copy_b.partition_S(gZ);        // (BCPY,BCPY_N,BCPY_K,k)
+  Tensor tBrZ = make_fragment_like(tBgZ(_,_,_,0)); // (BCPY,BCPY_N,BCPY_K)
+
+  // MMA.
+  ThrMMA thr_mma = mma.get_slice(thread_idx);
+  Tensor tCsA = thr_mma.partition_A(sA);       // (MMA,MMA_M,MMA_K)
+  Tensor tCsB = thr_mma.partition_B(sB);       // (MMA,MMA_N,MMA_K)
+  Tensor tCgC = thr_mma.partition_C(gC);       // (MMA,MMA_M,MMA_N)
+  Tensor tCrC = thr_mma.make_fragment_C(tCgC); // (MMA,MMA_M,MMA_N)
+
+  // Predicates for m/n bounds.
+  Tensor tApA = make_tensor<bool>(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{}); // (CPY_M,CPY_K)
+  Tensor tBpB = make_tensor<bool>(make_shape(size<1>(tBsB), size<2>(tBsB)), Stride<_1,_0>{}); // (CPY_N,CPY_K)
+  Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K)
+  Tensor cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB))); // (BLK_N,BLK_K)
+  Tensor cC = make_identity_tensor(make_shape(size<0>(gC), size<1>(gC))); // (M,N)
+  Tensor tAcA = thr_copy_a.partition_S(cA); // (CPY,CPY_M,CPY_K)
+  Tensor tBcB = thr_copy_b.partition_S(cB); // (CPY,CPY_N,CPY_K)
+  Tensor tCcC = thr_mma.partition_C(cC);    // (MMA,MMA_M,MMA_N)
+  CUTE_UNROLL
+  for (int m = 0; m < size<0>(tApA); ++m) {
+    tApA(m,0) = get<0>(tAcA(0,m,0)) < m_max_coord;
+  }
+  CUTE_UNROLL
+  for (int n = 0; n < size<0>(tBpB); ++n) {
+    tBpB(n,0) = get<0>(tBcB(0,n,0)) < n_max_coord;
+  }
+
+  // GMEM => RMEM.
+  auto fetch_gmem = [&](int tile) {
+    copy_if(copy_a, tApA, tAgA(_,_,_,tile), tArA);
+    copy_if(copy_b, tBpB, tBgB(_,_,_,tile), tBrB);
+    copy(tBgS(_,_,_,tile), tBrS);
+    copy(tBgZ(_,_,_,tile), tBrZ);
+  };
+  // RMEM => SMEM.
+  auto store_smem = [&]() {
+    __syncthreads();
+    copy(tArA, tAsA);
+    CUTE_UNROLL
+    for (int k = 0; k < size<2>(tBrB); ++k) {
+      CUTE_UNROLL
+      for (int n = 0; n < size<1>(tBrB); ++n) {
+        cute_dequant(tBrB(_,n,k), tBrS(_,n,k), tBrZ(_,n,k), tBrB_dq(_,n,k));
+      }
+    }
+    copy(tBrB_dq, tBsB);
+    __syncthreads();
+  };
+
+  // Clear the rmem tiles to account for predicated off loads.
+  if constexpr (HasKResidue) {
+    clear(tArA);
+    clear(tBrB);
+    clear(tBrS);
+    clear(tBrZ);
+  }
+
+  // Prefetch first tile.
+  if constexpr (HasKResidue) {
+    Tensor tAgA_k = tAgA(_,_,_,0);
+    CUTE_UNROLL
+    for (int k = 0; k < size<2>(tArA); ++k) {
+      if (get<1>(tAcA(0,0,k)) >= -k_residue) {
+        copy_if(copy_a, tApA(_,k), tAgA_k(_,_,k), tArA(_,_,k));
+      }
+    }
+    Tensor tBgB_k = tBgB(_,_,_,0);
+    Tensor tBgS_k = tBgS(_,_,_,0);
+    Tensor tBgZ_k = tBgZ(_,_,_,0);
+    CUTE_UNROLL
+    for (int k = 0; k < size<2>(tBrB); ++k) {
+      if (get<1>(tBcB(0,0,k)) >= -k_residue) {
+        copy_if(copy_b, tBpB(_,k), tBgB_k(_,_,k), tBrB(_,_,k));
+        copy(tBgS_k(_,_,k), tBrS(_,_,k));
+        copy(tBgZ_k(_,_,k), tBrZ(_,_,k));
+      }
+    }
+  } else {
+    fetch_gmem(0);
+  }
+
+  // Clear accumulators.
+  clear(tCrC);
+
+  // Loop over CTA tiles.
+  auto K_TILE_MAX = size<3>(tAgA);
+  for (int tile = 0; tile < K_TILE_MAX; ++tile) {
+    store_smem();
+    if constexpr (HasKResidue) {
+      // Avoid fetching full 0th-tile when there is residue.
+      if (K_TILE_MAX > 1) {
+        fetch_gmem((tile + 1 < K_TILE_MAX) ? tile + 1 : tile);
+      }
+    } else {
+      fetch_gmem((tile + 1 < K_TILE_MAX) ? tile + 1 : tile);
+    }
+    gemm(mma, tCsA, tCsB, tCrC);
+  }
+
+  // Epilogue.
+  CUTE_UNROLL
+  for (int i = 0; i < size(tCrC); ++i) {
+    if ((get<0>(tCcC(i)) < m_max_coord) && (get<1>(tCcC(i)) < n_max_coord)) {
+      tCgC(i) = Element(tCrC(i));
+    }
+  }
+}
+
+template <bool KMajor>
+inline constexpr auto make_matrix_stride(auto m, auto k) {
+  if constexpr (KMajor) {
+    return cute::make_stride(k, cute::Int<1>{}, m * k);
+  } else {
+    return cute::make_stride(cute::Int<1>{}, m, m * k);
+  }
+}
+
+template <bool KMajor>
+inline constexpr auto make_scales_layout(auto n, auto k, auto l, auto group_size) {
+  if constexpr (KMajor) {
+    return make_layout(
+        make_shape(n, make_shape(group_size, k / group_size), l),
+        make_stride(k / group_size, Stride<_0,_1>{}, n * k / group_size));
+  } else {
+    return make_layout(
+        make_shape(make_shape(group_size, n / group_size), k, l),
+        make_stride(Stride<_0,_1>{}, n / group_size, n * k / group_size));
+  }
+}
+
+template <int TileM, bool SM80>
+inline constexpr auto make_cta_tiler(auto group_size) {
+  auto bM = Int<TileM>{};
+  auto bN = Int<(!SM80 && group_size > 64) ? 64 : 128>{};
+  auto bK = Int<max(64, group_size)>{};
+  return make_shape(bM, bN, bK);
+}
+
+template <bool SM80, typename Element>
+inline constexpr auto make_tiled_mma(auto cta_tiler) {
+  using Atom = std::conditional_t<
+      SM80,
+      std::conditional_t<
+          std::is_same_v<Element, half_t>,
+          SM80_16x8x16_F32F16F16F32_TN,
+          std::conditional_t<
+              std::is_same_v<Element, bfloat16_t>,
+              SM80_16x8x16_F32BF16BF16F32_TN,
+              UniversalFMA<float>
+          >
+      >,
+      UniversalFMA<float, Element, Element>>;
+  if constexpr (!SM80 || std::is_same_v<Element, float>) {
+    return make_tiled_mma(Atom{}, Layout<Shape<_16,_8,_1>>{});
+  } else {
+    if constexpr (size<0>(cta_tiler) >= 32) {
+      return make_tiled_mma(Atom{}, Layout<Shape<_2,_2,_1>>{}, Tile<_32,_32,_16>{});
+    } else {
+      return make_tiled_mma(Atom{}, Layout<Shape<_1,_4,_1>>{}, Tile<_16,_32,_16>{});
+    }
+  }
+}
+
+} // namespace cutlass_gemm
+
+// clang-format on
+
+namespace mlx::core {
+
+template <typename F>
+inline void dispatch_element_types(Dtype dtype, const char* tag, F&& f) {
+  if (dtype == float32) {
+    f.template operator()<float>();
+  } else if (dtype == float16) {
+    f.template operator()<cutlass::half_t>();
+  } else if (dtype == bfloat16) {
+    f.template operator()<cutlass::bfloat16_t>();
+  } else {
+    throw std::invalid_argument(
+        fmt::format("{} Unsupported dtype: {}.", tag, dtype_to_string(dtype)));
+  }
+}
+
+template <typename F>
+inline void dispatch_groups(int group_size, const char* tag, F&& f) {
+  if (group_size == 32) {
+    f.template operator()<32>();
+  } else if (group_size == 64) {
+    f.template operator()<64>();
+  } else if (group_size == 128) {
+    f.template operator()<128>();
+  } else {
+    throw std::invalid_argument(
+        fmt::format("{} Group size {} is not supported.", tag, group_size));
+  }
+}
+
+template <typename T, typename F>
+inline void dispatch_quant_types(
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    const char* tag,
+    F&& f) {
+  if (mode == QuantizationMode::Mxfp4) {
+    f.template operator()<cutlass::float_e2m1_t, cutlass::float_ue8m0_t, 32>();
+  } else if (mode == QuantizationMode::Mxfp8) {
+    f.template operator()<cutlass::float_e4m3_t, cutlass::float_ue8m0_t, 32>();
+  } else if (mode == QuantizationMode::Nvfp4) {
+    f.template operator()<cutlass::float_e2m1_t, cutlass::float_e4m3_t, 16>();
+  } else {
+    dispatch_groups(group_size, tag, [&]<int group_size>() {
+      if (bits == 2) {
+        f.template operator()<cutlass::uint2b_t, T, group_size>();
+      } else if (bits == 3) {
+        f.template operator()<cutlass::uint3b_t, T, group_size>();
+      } else if (bits == 4) {
+        f.template operator()<cutlass::uint4b_t, T, group_size>();
+      } else if (bits == 5) {
+        f.template operator()<cutlass::uint5b_t, T, group_size>();
+      } else if (bits == 6) {
+        f.template operator()<cutlass::uint6b_t, T, group_size>();
+      } else if (bits == 8) {
+        f.template operator()<uint8_t, T, group_size>();
+      } else {
+        throw std::invalid_argument(
+            fmt::format("{} {}-bit quantization is not supported.", tag, bits));
+      }
+    });
+  }
+}
+
+} // namespace mlx::core

mlx/backend/cuda/quantized/qmm/qmm_sm80.cu

@@ -0,0 +1,240 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/quantized/qmm/qmm.h"
+#include "mlx/backend/cuda/quantized/qmm/qmm_sm80.cuh"
+
+// clang-format off
+
+// We can't put kernel code in mlx::core due to name conflicts of "Shape".
+namespace cutlass_gemm {
+
+using namespace cute;
+
+template <typename Element, typename Quant, typename Scale,
+          typename ProblemShape,
+          typename CtaTiler,
+          typename StrideA,
+          typename StrideB,
+          typename LayoutS,
+          typename StrideC,
+          typename TiledMma>
+__global__
+__launch_bounds__(decltype(size(TiledMma{}))::value)
+void qmm_sm80_kernel(
+    ProblemShape shape_MNKL, CtaTiler cta_tiler,
+    const Element* A, StrideA dA,
+    const Quant* B, StrideB dB,
+    const Scale* S, const Element* Z, LayoutS S_layout,
+    const uint32_t* lhs_indices, const uint32_t* rhs_indices,
+    Element* C, StrideC dC,
+    TiledMma mma) {
+  CUTE_STATIC_ASSERT_V(congruent(select<0,2,3>(shape_MNKL), dA));
+  CUTE_STATIC_ASSERT_V(congruent(select<1,2,3>(shape_MNKL), dB));
+  CUTE_STATIC_ASSERT_V(congruent(select<0,1,3>(shape_MNKL), dC));
+
+  int thread_idx = int(threadIdx.x);
+  int m_coord = int(blockIdx.x);
+  int n_coord = int(blockIdx.y);
+  int l_coord = int(blockIdx.z);
+
+  // For gather, use index lookup for input batch slicing.
+  uint32_t a_batch = lhs_indices ? lhs_indices[l_coord] : l_coord;
+  uint32_t b_batch = rhs_indices ? rhs_indices[l_coord] : l_coord;
+
+  // Represent the full tensors.
+  Tensor mA_mkl = make_tensor(make_gmem_ptr(A),        select<0,2,3>(shape_MNKL), dA); // (M,K,L)
+  Tensor mB_nkl = make_tensor(make_gmem_ptr<Quant>(B), select<1,2,3>(shape_MNKL), dB); // (N,K,L)
+  Tensor mC_mnl = make_tensor(make_gmem_ptr(C),        select<0,1,3>(shape_MNKL), dC); // (M,N,L)
+
+  Tensor mS_nkl = make_tensor(make_gmem_ptr(S), S_layout); // (N,(group_size,K/group_size),L)
+  Tensor mZ_nkl = make_tensor(make_gmem_ptr(Z), S_layout); // (N,(group_size,K/group_size),L)
+
+  // Get batch slice.
+  Tensor mA = mA_mkl(_,_,a_batch); // (M,K)
+  Tensor mB = mB_nkl(_,_,b_batch); // (N,K)
+  Tensor mC = mC_mnl(_,_,l_coord); // (M,N)
+
+  Tensor mS = mS_nkl(_,_,b_batch); // (N,(group_size,K/group_size))
+  Tensor mZ = mZ_nkl(_,_,b_batch); // (N,(group_size,K/group_size))
+
+  // Get the appropriate blocks for this thread block.
+  auto cta_coord = make_coord(m_coord, n_coord, _); // (m,n,k)
+  Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k)
+  Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
+  Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N)
+
+  Tensor gS = local_tile(mS, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
+  Tensor gZ = local_tile(mZ, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
+
+  // Compute tile residues for predication.
+  auto m_max_coord = size<0>(shape_MNKL) - size<0>(gA) * m_coord; // M - BLK_M * m_coord
+
+  qmm_sm80_mainloop(
+      cta_tiler,
+      gA,
+      gB,
+      gS,
+      gZ,
+      gC,
+      mma,
+      m_max_coord,
+      thread_idx);
+}
+
+template <int TileM,
+          typename Element, typename Quant, typename Scale>
+void qmm_sm80(
+    const Element* A,
+    const Quant*   B,
+    const Scale* S,
+    const Element* Z,
+    const uint32_t* lhs_indices,
+    const uint32_t* rhs_indices,
+    Element* C,
+    int m, int n, int k, int l,
+    bool broadcast_b,
+    auto group_size,
+    auto&& launch_kernel) {
+  // Define shapes (dynamic).
+  auto shape_MNKL = make_shape(m, n, k, l); // (M,N,K,L)
+
+  // Define layouts (mixed).
+  auto dA = make_stride(k, Int<1>{}, m * k); // (dM,dK,dL)
+  auto dB = make_stride(k, Int<1>{}, n * k); // (dN,dK,dL)
+  auto dC = make_stride(n, Int<1>{}, m * n); // (dM,dN,dL)
+  auto S_layout = make_scales_layout(n, k, l, group_size);
+
+  // Handle broadcasting.
+  if (broadcast_b) {
+    get<2>(dB) = 0;
+    get<2>(stride(S_layout)) = 0;
+  }
+
+  // Define CTA tile sizes (static).
+  auto cta_tiler = make_cta_tiler<TileM>(group_size);
+
+  // Define MMA.
+  TiledMMA mma = make_tiled_mma<TileM, Element>();
+  auto num_threads = size(mma);
+
+  // Shared memory size.
+  auto [sA_layout, sB_layout, sC_layout] = make_smem_layouts(cta_tiler);
+  size_t smem_bytes = sizeof(SharedStorage<Element, Quant,
+                                           decltype(sA_layout),
+                                           decltype(sB_layout),
+                                           decltype(sC_layout)>);
+
+  auto* kernel = &qmm_sm80_kernel<
+      Element, Quant, Scale,
+      decltype(shape_MNKL),
+      decltype(cta_tiler),
+      decltype(dA),
+      decltype(dB),
+      decltype(S_layout),
+      decltype(dC),
+      decltype(mma)>;
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
+
+  dim3 num_blocks{uint32_t(ceil_div(m, size<0>(cta_tiler))),
+                  uint32_t(ceil_div(n, size<1>(cta_tiler))),
+                  uint32_t(l)};
+  dim3 block_dims{num_threads};
+  void* args[] = {
+      &shape_MNKL, &cta_tiler,
+      &A, &dA,
+      &B, &dB,
+      &S, &Z, &S_layout,
+      &lhs_indices, &rhs_indices,
+      &C, &dC,
+      &mma};
+  launch_kernel(reinterpret_cast<void*>(kernel), num_blocks, block_dims, smem_bytes, args);
+}
+
+} // namespace cutlass_gemm
+
+// clang-format on
+
+namespace mlx::core {
+
+template <int TileM>
+void qmm_sm80_impl(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const std::optional<array>& lhs_indices,
+    const std::optional<array>& rhs_indices,
+    array& out,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::CommandEncoder& encoder) {
+  const char* tag = "[quantized_matmul]";
+  int m = out.ndim() > 1 ? out.shape(-2) : 1;
+  int n = out.shape(-1);
+  int k = x.shape(-1);
+  int l = out.size() / (m * n);
+  bool broadcast_b = (w.ndim() <= 2) || (w.size() != w.data_size());
+
+  dispatch_element_types(out.dtype(), tag, [&]<typename Element>() {
+    dispatch_quant_types<Element>(
+        bits,
+        group_size,
+        mode,
+        tag,
+        [&]<typename Quant, typename Scale, int group_size>() {
+          encoder.set_input_array(x);
+          encoder.set_input_array(w);
+          encoder.set_input_array(scales);
+          if (biases) {
+            encoder.set_input_array(*biases);
+          }
+          if (lhs_indices) {
+            encoder.set_input_array(*lhs_indices);
+          }
+          if (rhs_indices) {
+            encoder.set_input_array(*rhs_indices);
+          }
+          encoder.set_output_array(out);
+          cutlass_gemm::qmm_sm80<TileM>(
+              gpu_ptr<Element>(x),
+              gpu_ptr<Quant>(w),
+              gpu_ptr<Scale>(scales),
+              biases ? gpu_ptr<Element>(*biases) : nullptr,
+              lhs_indices ? gpu_ptr<uint32_t>(*lhs_indices) : nullptr,
+              rhs_indices ? gpu_ptr<uint32_t>(*rhs_indices) : nullptr,
+              gpu_ptr<Element>(out),
+              m,
+              n,
+              k,
+              l,
+              broadcast_b,
+              cute::Int<group_size>{},
+              [&](auto* kernel,
+                  dim3 num_blocks,
+                  dim3 block_dims,
+                  size_t smem_bytes,
+                  void** args) {
+                encoder.add_kernel_node_raw(
+                    kernel, num_blocks, block_dims, {}, smem_bytes, args);
+              });
+        });
+  });
+}
+
+// clang-format off
+template void qmm_sm80_impl<@TileM@>(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const std::optional<array>& biases,
+    const std::optional<array>& lhs_indices,
+    const std::optional<array>& rhs_indices,
+    array& out,
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    cu::CommandEncoder& encoder);
+// clang-format on
+
+} // namespace mlx::core

mlx/backend/cuda/quantized/qmm/qmm_sm80.cuh

@@ -0,0 +1,346 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/quantized/qmm/cute_dequant.cuh"
+#include "mlx/dtype_utils.h"
+
+// clang-format off
+
+// We can't put kernel code in mlx::core due to name conflicts of "Shape".
+namespace cutlass_gemm {
+
+using namespace cute;
+
+template <typename Element,
+          typename Quant,
+          typename SmemLayoutA,
+          typename SmemLayoutB,
+          typename SmemLayoutC>
+union SharedStorage {
+  struct {
+    ArrayEngine<Element, cosize_v<SmemLayoutA>> A;
+    ArrayEngine<Quant,   cosize_v<SmemLayoutB>> B;
+  } mainloop;
+  struct {
+    ArrayEngine<Element, cosize_v<SmemLayoutC>> C;
+  } epilogue;
+};
+
+inline constexpr auto make_smem_layouts(auto cta_tiler) {
+  // Define the A/B smem layouts (static).
+  auto swizzle_ab = composition(Swizzle<3,3,3>{},
+                                Layout<Shape <_8,Shape <_8, _8>>,
+                                       Stride<_8,Stride<_1,_64>>>{});
+  auto [bM, bN, bK] = cta_tiler;
+  auto bP = Int<3>{}; // pipeline
+  auto sA_layout = tile_to_shape(swizzle_ab, make_shape(bM, bK, bP));
+  auto sB_layout = tile_to_shape(swizzle_ab, make_shape(bN, bK, bP));
+
+  // Define the C smem layouts (static).
+  // TODO: Find a better swizzle.
+  auto sC_layout = tile_to_shape(swizzle_ab, make_shape(bM, bN));
+
+  return std::make_tuple(sA_layout, sB_layout, sC_layout);
+}
+
+template <typename T, int bits, template <typename U> typename Atom>
+inline constexpr auto make_tiled_copy(auto num_threads) {
+  return make_tiled_copy(
+      Copy_Atom<Atom<uint_bit_t<bits>>, T>{},
+      make_layout(make_shape(Int<num_threads / 8>{}, Int<8>{}), LayoutRight{}),
+      make_layout(make_shape(Int<1>{}, Int<bits / sizeof_bits_v<T>>{})));
+}
+
+template <typename CtaTiler,
+          typename TensorA,
+          typename TensorB,
+          typename TensorS,
+          typename TensorZ,
+          typename TensorC,
+          typename TiledMma>
+CUTE_DEVICE void qmm_sm80_mainloop(
+    CtaTiler cta_tiler,
+    TensorA gA,
+    TensorB gB,
+    TensorS gS,
+    TensorZ gZ,
+    TensorC gC,
+    TiledMma mma,
+    int m_max_coord,
+    int thread_idx) {
+  // Get the types of operands.
+  using Element = decltype(gA)::value_type;
+  using Quant = decltype(gB)::value_type;
+  using Scale = decltype(gS)::value_type;
+
+  // Define smem layouts.
+  auto [sA_layout, sB_layout, sC_layout] = make_smem_layouts(cta_tiler);
+
+  // Shared memory buffer.
+  extern __shared__ char smem_buf[];
+  using SharedStorage = SharedStorage<Element, Quant,
+                                      decltype(sA_layout),
+                                      decltype(sB_layout),
+                                      decltype(sC_layout)>;
+  SharedStorage& smem = *reinterpret_cast<SharedStorage*>(smem_buf);
+  Tensor sA = make_tensor(make_smem_ptr(smem.mainloop.A.begin()), sA_layout); // (BLK_M,BLK_K)
+  Tensor sB = make_tensor(make_smem_ptr(smem.mainloop.B.begin()), sB_layout); // (BLK_N,BLK_K)
+  Tensor sC = make_tensor(make_smem_ptr(smem.epilogue.C.begin()), sC_layout); // (BLK_M,BLK_N)
+
+  // Define copy atoms.
+  constexpr int element_bits = sizeof_bits_v<Element>;
+  constexpr int quant_bits = sizeof_bits_v<Quant>;
+  constexpr int qload = 128 / (element_bits / quant_bits);
+  auto num_threads = size(mma);
+  TiledCopy g2s_copy_a = make_tiled_copy<Element, 128, SM80_CP_ASYNC_CACHEALWAYS>(num_threads);
+  TiledCopy g2s_copy_b = make_tiled_copy<Quant, qload, SM80_CP_ASYNC_CACHEALWAYS>(num_threads);
+  TiledCopy s2g_copy_c = make_tiled_copy<Element, 128, UniversalCopy>(num_threads);
+
+  Copy_Atom<SM75_U32x4_LDSM_N, Element> s2r_atom_a;
+  Copy_Atom<UniversalCopy<uint_bit_t<2 * quant_bits>>, Quant> s2r_atom_b;
+  Copy_Atom<UniversalCopy<uint_bit_t<2 * element_bits>>, Element> r2s_atom_c;
+  Copy_Atom<UniversalCopy<Scale>, Scale> g2r_atom_s;
+
+  // Partition the copying of A/B/C tiles across the threads.
+  ThrCopy g2s_thr_copy_a = g2s_copy_a.get_slice(thread_idx);
+  Tensor tAgA = g2s_thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k)
+  Tensor tAsA = g2s_thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K,PIPE)
+
+  ThrCopy g2s_thr_copy_b = g2s_copy_b.get_slice(thread_idx);
+  Tensor tBgB = g2s_thr_copy_b.partition_S(gB);  // (BCPY,BCPY_N,BCPY_K,k)
+  Tensor tBsB = g2s_thr_copy_b.partition_D(sB);  // (BCPY,BCPY_N,BCPY_K,PIPE)
+
+  ThrCopy s2g_thr_copy_c = s2g_copy_c.get_slice(thread_idx);
+  Tensor s2g_tCsC = s2g_thr_copy_c.partition_S(sC); // (CCPY,CCPY_M,CCPY_N)
+  Tensor s2g_tCgC = s2g_thr_copy_c.partition_D(gC); // (CCPY,CCPY_M,CCPY_N)
+
+  // MMA.
+  ThrMMA thr_mma = mma.get_slice(thread_idx);
+  Tensor tCrA = thr_mma.partition_fragment_A(sA(_,_,0)); // (MMA,MMA_M,MMA_K)
+  Tensor tCsB = thr_mma.partition_B(sB(_,_,0));          // (MMA,MMA_N,MMA_K)
+  Tensor tCrB = make_fragment_like<Quant>(tCsB);         // (MMA,MMA_N,MMA_K)
+  Tensor tCrB_dq = make_fragment_like<Element>(tCsB);    // (MMA,MMA_N,MMA_K)
+  Tensor tCgC = thr_mma.partition_C(gC);                 // (MMA,MMA_M,MMA_N)
+  Tensor tCrC_accu = make_fragment_like<float>(tCgC);    // (MMA,MMA_M,MMA_N)
+  Tensor tCrC = make_fragment_like<Element>(tCgC);       // (MMA,MMA_M,MMA_N)
+
+  Tensor tCgS = thr_mma.partition_B(gS);         // (MMA,MMA_N,MMA_K,k)
+  Tensor tCrS = make_tensor_like(tCgS(_,_,_,0)); // (MMA,MMA_N,MMA_K)
+  Tensor tCgZ = thr_mma.partition_B(gZ);         // (MMA,MMA_N,MMA_K,k)
+  Tensor tCrZ = make_tensor_like(tCgZ(_,_,_,0)); // (MMA,MMA_N,MMA_K)
+
+  // Copy Atom retiling.
+  TiledCopy s2r_copy_a = make_tiled_copy_A(s2r_atom_a, mma);
+  ThrCopy s2r_thr_copy_a = s2r_copy_a.get_slice(thread_idx);
+  Tensor s2r_tCsA = s2r_thr_copy_a.partition_S(sA); // (ACPY,MMA_M,MMA_K,PIPE)
+  Tensor s2r_tCrA = s2r_thr_copy_a.retile_D(tCrA);  // (ACPY,MMA_M,MMA_K)
+
+  TiledCopy s2r_copy_b = make_tiled_copy_B(s2r_atom_b, mma);
+  ThrCopy s2r_thr_copy_b = s2r_copy_b.get_slice(thread_idx);
+  Tensor s2r_tCsB = s2r_thr_copy_b.partition_S(sB); // (BCPY,MMA_N,MMA_K,PIPE)
+  Tensor s2r_tCrB = s2r_thr_copy_b.retile_D(tCrB);  // (BCPY,MMA_N,MMA_K)
+
+  TiledCopy r2s_copy_c = make_tiled_copy_C(r2s_atom_c, mma);
+  ThrCopy r2s_thr_copy_c = r2s_copy_c.get_slice(thread_idx);
+  Tensor r2s_tCrC = r2s_thr_copy_c.retile_S(tCrC);  // (CCPY,MMA_M,MMA_N)
+  Tensor r2s_tCsC = r2s_thr_copy_c.partition_D(sC); // (CCPY,MMA_M,MMA_N)
+
+  TiledCopy g2r_copy_s = make_tiled_copy_B(g2r_atom_s, mma);
+  ThrCopy g2r_thr_copy_s = g2r_copy_s.get_slice(thread_idx);
+  Tensor g2r_tCgS = g2r_thr_copy_s.partition_S(gS); // (BCPY,MMA_N,MMA_K,k)
+  Tensor g2r_tCrS = g2r_thr_copy_s.retile_D(tCrS);  // (BCPY,MMA_N,MMA_K)
+  Tensor g2r_tCgZ = g2r_thr_copy_s.partition_S(gZ); // (BCPY,MMA_N,MMA_K,k)
+  Tensor g2r_tCrZ = g2r_thr_copy_s.retile_D(tCrZ);  // (BCPY,MMA_N,MMA_K)
+
+  // Predicates for m bound.
+  Tensor tApA = make_tensor<bool>(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{});         // (CPY_M,CPY_K)
+  Tensor tCpC = make_tensor<bool>(make_shape(size<1>(s2g_tCsC), size<2>(s2g_tCsC)), Stride<_1,_0>{}); // (CPY_M,CPY_N)
+  Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K)
+  Tensor cC = make_identity_tensor(make_shape(size<0>(sC), size<1>(sC))); // (BLK_M,BLK_N)
+  Tensor tAcA = g2s_thr_copy_a.partition_D(cA); // (CPY,CPY_M,CPY_K)
+  Tensor tCcC = s2g_thr_copy_c.partition_D(cC); // (CPY,CPY_M,CPY_N)
+  CUTE_UNROLL
+  for (int m = 0; m < size<0>(tApA); ++m) {
+    tApA(m,0) = get<0>(tAcA(0,m,0)) < m_max_coord;
+  }
+  CUTE_UNROLL
+  for (int m = 0; m < size<0>(tCpC); ++m) {
+    tCpC(m,0) = get<0>(tCcC(0,m,0)) < m_max_coord;
+  }
+
+  auto K_PIPE_MAX = size<3>(tAsA);
+  int smem_pipe_read = 0;
+  int smem_pipe_write = 0;
+
+  // Copy A/B: GMEM => SMEM.
+  auto fetch_gmem = [&](int tile) {
+    copy_if(g2s_copy_a, tApA, tAgA(_,_,_,tile), tAsA(_,_,_,smem_pipe_write));
+    copy(g2s_copy_b, tBgB(_,_,_,tile), tBsB(_,_,_,smem_pipe_write));
+    cp_async_fence();
+    smem_pipe_write = (smem_pipe_write + 1) % K_PIPE_MAX;
+  };
+  // Copy S/Z: GMEM => RMEM.
+  auto fetch_scales = [&](int tile) {
+    copy(g2r_copy_s, g2r_tCgS(_,_,_,tile), g2r_tCrS);
+    if constexpr (quant_has_bias_v<Quant>) {
+      copy(g2r_copy_s, g2r_tCgZ(_,_,_,tile), g2r_tCrZ);
+    }
+  };
+  // Copy A/B: SMEM => RMEM.
+  auto fetch_smem = [&](auto block) {
+    copy(s2r_atom_a, s2r_tCsA(_,_,block,smem_pipe_read), s2r_tCrA(_,_,block));
+    copy(s2r_atom_b, s2r_tCsB(_,_,block,smem_pipe_read), s2r_tCrB(_,_,block));
+    CUTE_UNROLL
+    for (int n = 0; n < size<1>(tCrB); ++n) {
+      cute_vectorized_dequant(
+          tCrB(_,n,block),
+          tCrS(_,n,block),
+          tCrZ(_,n,block),
+          tCrB_dq(_,n,block));
+    }
+  };
+
+  auto K_TILE_MAX = size<3>(tAgA);
+  auto K_BLOCK_MAX = size<2>(tCrA);
+
+  // Prefetch beginning tiles.
+  int tile_pipe = 0;
+  CUTE_UNROLL
+  for (; tile_pipe < K_PIPE_MAX - 1; ++tile_pipe) {
+    fetch_gmem(tile_pipe);
+  }
+
+  // Clear accumulators.
+  clear(tCrC_accu);
+
+  // Prefetch first block.
+  if constexpr (K_BLOCK_MAX > 1) {
+    cp_async_wait<K_PIPE_MAX - 2>();
+    __syncthreads();
+    fetch_scales(0);
+    fetch_smem(Int<0>{});
+  }
+
+  // Loop over CTA tiles.
+  for (int tile = 0; tile < K_TILE_MAX; ++tile) {
+    // Unroll MMA blocks.
+    CUTE_UNROLL
+    for (int block = 0; block < K_BLOCK_MAX; ++block) {
+      // Wait for last tile.
+      if (block == K_BLOCK_MAX - 1) {
+        smem_pipe_read = (smem_pipe_read + 1) % K_PIPE_MAX;
+        cp_async_wait<K_PIPE_MAX - 2>();
+        __syncthreads();
+        fetch_scales((tile + 1 < K_TILE_MAX) ? tile + 1 : tile);
+      }
+      // Prefetch next block.
+      fetch_smem((block + 1) % K_BLOCK_MAX);
+      // Prefetch next tile.
+      if (block == 0) {
+        fetch_gmem(tile_pipe);
+        tile_pipe = (tile_pipe + 1 < K_TILE_MAX) ? tile_pipe + 1 : tile_pipe;
+      }
+      // MMA.
+      gemm(mma, tCrA(_,_,block), tCrB_dq(_,_,block), tCrC_accu);
+    }
+  }
+
+  // Epilogue.
+  CUTE_UNROLL
+  for (int i = 0; i < size(tCrC_accu); i++) {
+    tCrC(i) = Element(tCrC_accu(i));
+  }
+  copy(r2s_copy_c, r2s_tCrC, r2s_tCsC);
+  __syncthreads();
+  copy_if(s2g_copy_c, tCpC, s2g_tCsC, s2g_tCgC);
+}
+
+inline constexpr auto make_scales_layout(auto n, auto k, auto l, auto group_size) {
+  return make_layout(
+      make_shape(n, make_shape(group_size, k / group_size), l),
+      make_stride(k / group_size, Stride<_0,_1>{}, n * k / group_size));
+}
+
+template <int TileM>
+inline constexpr auto make_cta_tiler(auto group_size) {
+  auto bM = Int<TileM>{};
+  auto bN = Int<128>{};
+  auto bK = Int<max(64, group_size)>{};
+  return make_shape(bM, bN, bK);
+}
+
+template <int TileM, typename Element>
+inline constexpr auto make_tiled_mma() {
+  using Atom = std::conditional_t<
+      std::is_same_v<Element, half_t>,
+      SM80_16x8x16_F32F16F16F32_TN,
+      std::conditional_t<
+          std::is_same_v<Element, bfloat16_t>,
+          SM80_16x8x16_F32BF16BF16F32_TN,
+          UniversalFMA<float>>>;
+  if constexpr (TileM >= 32) {
+    return make_tiled_mma(Atom{}, Layout<Shape<_2,_2,_1>>{}, Tile<_32,_32,_16>{});
+  } else {
+    return make_tiled_mma(Atom{}, Layout<Shape<_1,_4,_1>>{}, Tile<_16,_32,_16>{});
+  }
+}
+
+} // namespace cutlass_gemm
+
+// clang-format on
+
+namespace mlx::core {
+
+template <typename F>
+inline void dispatch_element_types(Dtype dtype, const char* tag, F&& f) {
+  if (dtype == float16) {
+    f.template operator()<cutlass::half_t>();
+  } else if (dtype == bfloat16) {
+    f.template operator()<cutlass::bfloat16_t>();
+  } else {
+    throw std::invalid_argument(
+        fmt::format("{} Unsupported dtype: {}.", tag, dtype_to_string(dtype)));
+  }
+}
+
+template <typename F>
+inline void dispatch_groups(int group_size, const char* tag, F&& f) {
+  if (group_size == 32) {
+    f.template operator()<32>();
+  } else if (group_size == 64) {
+    f.template operator()<64>();
+  } else if (group_size == 128) {
+    f.template operator()<128>();
+  } else {
+    throw std::invalid_argument(
+        fmt::format("{} Group size {} is not supported.", tag, group_size));
+  }
+}
+
+template <typename T, typename F>
+inline void dispatch_quant_types(
+    int bits,
+    int group_size,
+    QuantizationMode mode,
+    const char* tag,
+    F&& f) {
+  if (mode == QuantizationMode::Mxfp4) {
+    f.template operator()<cutlass::float_e2m1_t, cutlass::float_ue8m0_t, 32>();
+  } else if (mode == QuantizationMode::Mxfp8) {
+    f.template operator()<cutlass::float_e4m3_t, cutlass::float_ue8m0_t, 32>();
+  } else if (mode == QuantizationMode::Nvfp4) {
+    f.template operator()<cutlass::float_e2m1_t, cutlass::float_e4m3_t, 16>();
+  } else {
+    dispatch_groups(group_size, tag, [&]<int group_size>() {
+      if (bits == 4) {
+        f.template operator()<cutlass::uint4b_t, T, group_size>();
+      } else if (bits == 8) {
+        f.template operator()<uint8_t, T, group_size>();
+      } else {
+        throw std::invalid_argument(
+            fmt::format("{} {}-bit quantization is not supported.", tag, bits));
+      }
+    });
+  }
+}
+
+} // namespace mlx::core

mlx/backend/cuda/quantized/qmm/qmm_sm90.cu

@@ -20,8 +20,7 @@ namespace cutlass_gemm {
 using namespace cute;
 
 template <
-    typename TileShapeMN = Shape<_128, _16>,
-    typename ClusterShape = Shape<_1, _1, _1>,
+    int TileN = 16,
     typename Element,
     typename Quant,
     typename GroupSize,
@@ -36,6 +35,7 @@ void qmm_sm90(
     int64_t n,
     int64_t k,
     int64_t l,
+    bool broadcast_b,
     GroupSize group_size,
     F&& launch_kernel) {
   constexpr int kAlignmentA = 128 / sizeof_bits<Element>::value;
@@ -46,7 +46,8 @@ void qmm_sm90(
 
   using Arch = cutlass::arch::Sm90;
   using Accumulator = float;
-  using TileShape = decltype(append(TileShapeMN{}, Int<kTileShapeK>{}));
+  using TileShape = Shape<_128, Int<TileN>, Int<kTileShapeK>>;
+  using ClusterShape = Shape<Int<(TileN <= 32) ? 1 : 2>, _1, _1>;
 
   using Epilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
       Arch,
@@ -93,6 +94,10 @@ void qmm_sm90(
   auto dB = make_stride(k, Int<1>{}, n * k);
   auto dS = make_stride(Int<1>{}, n, n * k / group_size);
   auto dD = make_stride(Int<1>{}, n, m * n);
+  if (broadcast_b) {
+    get<2>(dB) = 0;
+    get<2>(dS) = 0;
+  }
 
   Gemm gemm;
   typename Gemm::Arguments args{
@@ -111,7 +116,9 @@ void qmm_sm90(
       reinterpret_cast<void*>(kernel),
       gemm.get_grid_shape(gemm.params()),
       GemmKernel::get_block_shape(),
-      {get<0>(cluster), get<1>(cluster), get<2>(cluster)},
+      {static_cast<unsigned>(get<0>(cluster)),
+       static_cast<unsigned>(get<1>(cluster)),
+       static_cast<unsigned>(get<2>(cluster))},
       GemmKernel::SharedStorageSize,
       kernel_params);
 }
@@ -170,8 +177,8 @@ inline void dispatch_groups(int group_size, const char* tag, F&& f) {
   }
 }
 
-template <typename TileShapeMN, typename ClusterShape>
-void qmm_impl_sm90(
+template <int TileN>
+void qmm_sm90_impl(
     const array& x,
     const array& w,
     const array& scales_,
@@ -182,29 +189,11 @@ void qmm_impl_sm90(
     cu::CommandEncoder& encoder,
     Stream s) {
   const char* tag = "[quantized_matmul]";
-  int m = out.shape(-2);
+  int m = out.ndim() > 1 ? out.shape(-2) : 1;
   int n = out.shape(-1);
   int k = x.shape(-1);
   int l = out.size() / (m * n);
-  if (k % 64 != 0) {
-    throw std::runtime_error(fmt::format("{} K must be multiples of 64.", tag));
-  }
-  if (!x.flags().row_contiguous) {
-    throw std::runtime_error(
-        fmt::format("{} Activations must be row contiguous.", tag));
-  }
-  if (!w.flags().row_contiguous) {
-    throw std::runtime_error(
-        fmt::format("{} Weights must be row contiguous.", tag));
-  }
-  if (!scales_.flags().row_contiguous) {
-    throw std::runtime_error(
-        fmt::format("{} Scales must be row contiguous.", tag));
-  }
-  if (!biases_.flags().row_contiguous) {
-    throw std::runtime_error(
-        fmt::format("{} Biases must be row contiguous.", tag));
-  }
+  bool broadcast_b = (w.ndim() <= 2) || (w.size() != w.data_size());
 
   // FIXME: Copy happens for every call.
   array scales = transpose_last_2_dims(scales_, encoder, s);
@@ -218,7 +207,7 @@ void qmm_impl_sm90(
         encoder.set_input_array(scales);
         encoder.set_input_array(biases);
         encoder.set_output_array(out);
-        cutlass_gemm::qmm_sm90(
+        cutlass_gemm::qmm_sm90<TileN>(
             gpu_ptr<Element>(x),
             gpu_ptr<Quant>(w),
             gpu_ptr<Element>(scales),
@@ -228,6 +217,7 @@ void qmm_impl_sm90(
             n,
             k,
             l,
+            broadcast_b,
             group_size,
             [&](auto* kernel,
                 dim3 num_blocks,
@@ -248,24 +238,19 @@ void qmm_impl_sm90(
   });
 }
 
+// clang-format off
+template void qmm_sm90_impl<@TileN@>(
+    const array& x,
+    const array& w,
+    const array& scales,
+    const array& biases,
+    array& out,
+    int bits,
+    int group_size,
+    cu::CommandEncoder& encoder,
+    Stream s);
+// clang-format on
+
 } // namespace mlx::core
 
-#define QMM_SM90_GPU(TileShapeMN, ClusterShape)           \
-  namespace mlx::core {                                   \
-  template void qmm_impl_sm90<TileShapeMN, ClusterShape>( \
-      const array& x,                                     \
-      const array& w,                                     \
-      const array& scales,                                \
-      const array& biases,                                \
-      array& out,                                         \
-      int bits,                                           \
-      int group_size,                                     \
-      cu::CommandEncoder& encoder,                        \
-      Stream s);                                          \
-  }
-
-#else
-
-#define QMM_SM90_GPU(TileShapeMN, ClusterShape)
-
 #endif // defined(MLX_CUDA_SM90A_ENABLED)