[CUDA] Vectorize generated kernels (#2444)

* Add more nvtx loggings

* [CUDA] Saving primitive inputs faster

* Remove unneeded check

.circleci/config.yml

@@ -81,23 +81,24 @@ jobs:
             export DEBIAN_FRONTEND=noninteractive
             export NEEDRESTART_MODE=a
             sudo apt-get update
-            sudo apt-get upgrade -y
-            pip install --upgrade cmake
             sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev
             sudo apt-get install openmpi-bin openmpi-common libopenmpi-dev
+            curl -LsSf https://astral.sh/uv/install.sh | sh
       - run:
           name: Install Python package
           command: |
-            pip install -e ".[dev]"
+            uv venv
+            uv pip install cmake
+            uv pip install -e ".[dev]" -v
       - run:
           name: Generate package stubs
           command: |
-            echo "stubs"
-            pip install typing_extensions
-            python setup.py generate_stubs
+            uv pip install typing_extensions
+            uv run --no-project setup.py generate_stubs
       - run:
           name: Run Python tests
           command: |
+            source .venv/bin/activate
             python -m unittest discover python/tests -v
             mpirun --bind-to none -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py
             mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2)
@@ -105,6 +106,7 @@ jobs:
       - run:
           name: Build CPP only
           command: |
+            source .venv/bin/activate
             mkdir -p build && cd build
             cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG
             make -j `nproc`
@@ -130,33 +132,30 @@ jobs:
       - run:
           name: Install dependencies
           command: |
-            brew install python@3.9
-            brew install openmpi
-            python3.9 -m venv env
-            source env/bin/activate
-            pip install --upgrade pip
-            pip install --upgrade cmake
-            pip install nanobind==2.4.0
-            pip install numpy
-            pip install torch
-            pip install tensorflow
-            pip install unittest-xml-reporting
+            HOMEBREW_NO_AUTO_UPDATE=1 HOMEBREW_NO_INSTALL_CLEANUP=1 \
+              brew install openmpi uv
       - run:
           name: Install Python package
           command: |
-            source env/bin/activate
+            uv venv --python 3.9
+            uv pip install \
+              nanobind==2.4.0 \
+              cmake \
+              numpy \
+              torch \
+              tensorflow \
+              unittest-xml-reporting
             DEBUG=1 CMAKE_ARGS="-DCMAKE_COMPILE_WARNING_AS_ERROR=ON" \
-              pip install -e . -v
+              uv pip install -e . -v
       - run:
           name: Generate package stubs
           command: |
-            source env/bin/activate
-            pip install typing_extensions
-            python setup.py generate_stubs
+            uv pip install typing_extensions
+            uv run --no-project setup.py generate_stubs
       - run:
           name: Run Python tests
           command: |
-            source env/bin/activate
+            source .venv/bin/activate
             LOW_MEMORY=1 DEVICE=cpu python -m xmlrunner discover -v python/tests -o test-results/cpu
             LOW_MEMORY=1 DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -m xmlrunner discover -v python/tests -o test-results/gpu
             mpirun --bind-to none -host localhost:8 -np 8 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python python/tests/mpi_test_distributed.py
@@ -165,16 +164,17 @@ jobs:
       - run:
           name: Build example extension
           command: |
-            source env/bin/activate
+            source .venv/bin/activate
             cd examples/extensions
-            pip install -r requirements.txt
-            python setup.py build_ext -j8
+            uv pip install -r requirements.txt
+            uv run --no-project setup.py build_ext --inplace
+            uv run --no-project python test.py
       - store_test_results:
           path: test-results
       - run:
           name: Build CPP only
           command: |
-            source env/bin/activate
+            source .venv/bin/activate
             mkdir -p build && cd build && cmake .. && make -j `sysctl -n hw.ncpu`
       - run:
           name: Run CPP tests
@@ -183,7 +183,7 @@ jobs:
       - run:
           name: Build small binary
           command: |
-            source env/bin/activate
+            source .venv/bin/activate
             cd build/
             cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \
               -DBUILD_SHARED_LIBS=ON \
@@ -195,12 +195,13 @@ jobs:
       - run:
           name: Run Python tests with JIT
           command: |
-            source env/bin/activate
             CMAKE_ARGS="-DMLX_METAL_JIT=ON" \
-              pip install -e . -v
+              uv pip install -e .
             LOW_MEMORY=1 DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 \
               METAL_DEBUG_ERROR_MODE=0 \
-              python -m xmlrunner discover -v python/tests -o test-results/gpu_jit
+              uv run --no-project python -m xmlrunner discover \
+                -v python/tests \
+                -o test-results/gpu_jit
 
   cuda_build_and_test:
     parameters:
@@ -212,22 +213,42 @@ jobs:
       resource_class: gpu.nvidia.small.gen2
     steps:
       - checkout
+      - restore_cache:
+          keys:
+            - cuda-<< parameters.image_date >>-{{ arch }}-
       - run:
-          name: Install Python package
+          name: Install dependencies
           command: |
             sudo apt-get update
             sudo apt-get install libcudnn9-dev-cuda-12
             sudo apt-get install libblas-dev liblapack-dev liblapacke-dev
-            python3 -m venv env
-            source env/bin/activate
+            curl -sL https://github.com/ccache/ccache/releases/download/v4.11.3/ccache-4.11.3-linux-x86_64.tar.xz | tar xJf -
+            sudo mv ccache-4.11.3-linux-x86_64/ccache /usr/bin/ccache
+            rm -rf ccache-4.11.3-linux-x86_64
+            curl -LsSf https://astral.sh/uv/install.sh | sh
+      - run:
+          name: Install Python package
+          command: |
+            uv venv
             CMAKE_ARGS="-DMLX_BUILD_CUDA=ON -DCMAKE_CUDA_COMPILER=`which nvcc`" \
-              pip install -e ".[dev]"
+              uv pip install -e ".[dev]" -v
       - run:
           name: Run Python tests
           command: |
-            source env/bin/activate
+            source .venv/bin/activate
             LOW_MEMORY=1 DEVICE=cpu python -m unittest discover python/tests -v
             LOW_MEMORY=1 DEVICE=gpu python -m tests discover python/tests -v
+      - run:
+          name: CCache report
+          command: |
+            ccache --show-stats
+            ccache --zero-stats
+            ccache --max-size 400MB
+            ccache --cleanup
+      - save_cache:
+          key: cuda-<< parameters.image_date >>-{{ arch }}-{{ epoch }}
+          paths:
+            - /home/circleci/.cache/ccache
 
   build_release:
     parameters:
@@ -323,14 +344,10 @@ jobs:
             export DEBIAN_FRONTEND=noninteractive
             export NEEDRESTART_MODE=a
             sudo apt-get update
-            sudo apt-get upgrade -y
             TZ=Etc/UTC sudo apt-get -y install tzdata
-            sudo apt-get install -y apt-utils
-            sudo apt-get install -y software-properties-common
             sudo add-apt-repository -y ppa:deadsnakes/ppa
             sudo apt-get install -y $PYTHON $PYTHON-dev $PYTHON-full
             sudo apt-get install -y libblas-dev liblapack-dev liblapacke-dev
-            sudo apt-get install -y build-essential git
             $PYTHON -m venv env
             source env/bin/activate
             pip install --upgrade pip
@@ -555,6 +572,9 @@ workflows:
           requires: [ hold ]
       - cuda_build_and_test:
           requires: [ hold ]
+          matrix:
+            parameters:
+              image_date: ["2023.11.1", "2025.05.1"]
   nightly_build:
     when:
       and:

CMakeLists.txt

@@ -41,6 +41,7 @@ option(MLX_BUILD_GGUF "Include support for GGUF format" ON)
 option(MLX_BUILD_SAFETENSORS "Include support for safetensors format" ON)
 option(MLX_BUILD_BLAS_FROM_SOURCE "Build OpenBLAS from source code" OFF)
 option(MLX_METAL_JIT "Use JIT compilation for Metal kernels" OFF)
+option(MLX_USE_CCACHE "Use CCache for compilation cache when available" ON)
 option(BUILD_SHARED_LIBS "Build mlx as a shared library" OFF)
 
 # --------------------- Processor tests -------------------------
@@ -68,6 +69,15 @@ else()
   set(MLX_BUILD_METAL OFF)
 endif()
 
+if(MLX_USE_CCACHE)
+  find_program(CCACHE_PROGRAM ccache)
+  if(CCACHE_PROGRAM)
+    set(CMAKE_C_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
+    set(CMAKE_CXX_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
+    set(CMAKE_CUDA_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
+  endif()
+endif()
+
 # ----------------------------- Lib -----------------------------
 
 include(FetchContent)

docs/src/dev/extensions.rst

@@ -394,14 +394,14 @@ below.
         out.set_data(allocator::malloc(out.nbytes()));
 
         // Resolve name of kernel
-        std::ostringstream kname;
-        kname << "axpby_" << "general_" << type_to_name(out);
+        std::stream kname;
+        kname = "axpby_general_" + type_to_name(out);
 
         // Load the metal library
-        auto lib = d.get_library("mlx_ext");
+        auto lib = d.get_library("mlx_ext", current_binary_dir());
 
         // Make a kernel from this metal library
-        auto kernel = d.get_kernel(kname.str(), lib);
+        auto kernel = d.get_kernel(kname, lib);
 
         // Prepare to encode kernel
         auto& compute_encoder = d.get_command_encoder(s.index);

examples/extensions/axpby/axpby.cpp

@@ -1,5 +1,6 @@
 // Copyright © 2023-2025 Apple Inc.
 
+#include <dlfcn.h>
 #include <iostream>
 #include <sstream>
 
@@ -16,6 +17,19 @@
 
 namespace my_ext {
 
+// A helper function to find the location of the current binary on disk.
+// The Metal library ("mlx_ext.mtllib"), should be in the same directory.
+std::string current_binary_dir() {
+  static std::string binary_dir = []() {
+    Dl_info info;
+    if (!dladdr(reinterpret_cast<void*>(&current_binary_dir), &info)) {
+      throw std::runtime_error("Unable to get current binary dir.");
+    }
+    return std::filesystem::path(info.dli_fname).parent_path().string();
+  }();
+  return binary_dir;
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 // Operation Implementation
 ///////////////////////////////////////////////////////////////////////////////
@@ -167,16 +181,15 @@ void Axpby::eval_gpu(
   }
 
   // Resolve name of kernel (corresponds to axpby.metal)
-  std::ostringstream kname;
-  kname << "axpby_";
-  kname << (contiguous_kernel ? "contiguous_" : "general_");
-  kname << type_to_name(out);
+  std::string kname = "axpby_";
+  kname += (contiguous_kernel ? "contiguous_" : "general_");
+  kname += type_to_name(out);
 
   // Load the metal library
-  auto lib = d.get_library("mlx_ext");
+  auto lib = d.get_library("mlx_ext", current_binary_dir());
 
   // Make a kernel from this metal library
-  auto kernel = d.get_kernel(kname.str(), lib);
+  auto kernel = d.get_kernel(kname, lib);
 
   // Prepare to encode kernel
   auto& compute_encoder = d.get_command_encoder(s.index);

examples/extensions/requirements.txt

@@ -1,4 +1,4 @@
 setuptools>=42
 cmake>=3.25
 mlx>=0.21.0
-nanobind==2.2.0
+nanobind==2.4.0

examples/extensions/test.py

@@ -3,8 +3,10 @@ from mlx_sample_extensions import axpby
 
 a = mx.ones((3, 4))
 b = mx.ones((3, 4))
-c = axpby(a, b, 4.0, 2.0, stream=mx.cpu)
+c_cpu = axpby(a, b, 4.0, 2.0, stream=mx.cpu)
+c_gpu = axpby(a, b, 4.0, 2.0, stream=mx.gpu)
 
-print(f"c shape: {c.shape}")
-print(f"c dtype: {c.dtype}")
-print(f"c correct: {mx.all(c == 6.0).item()}")
+print(f"c shape: {c_cpu.shape}")
+print(f"c dtype: {c_cpu.dtype}")
+print(f"c_cpu correct: {mx.all(c_cpu == 6.0).item()}")
+print(f"c_gpu correct: {mx.all(c_gpu == 6.0).item()}")

mlx/backend/common/utils.cpp

@@ -228,4 +228,31 @@ std::pair<Dims, Dims> get_grid_and_block_common(int dim0, int dim1, int dim2) {
       std::make_tuple(gx, gy, gz), std::make_tuple(bx, by, bz));
 }
 
+array swapaxes_in_eval(const array& x, int axis1, int axis2) {
+  int ndim = x.ndim();
+  if (axis1 < 0) {
+    axis1 += ndim;
+  }
+  if (axis2 < 0) {
+    axis2 += ndim;
+  }
+
+  auto shape = x.shape();
+  std::swap(shape[axis1], shape[axis2]);
+  auto strides = x.strides();
+  std::swap(strides[axis1], strides[axis2]);
+
+  auto [data_size, row_contiguous, col_contiguous] =
+      check_contiguity(shape, strides);
+  bool contiguous = data_size == x.data_size();
+
+  array out(std::move(shape), x.dtype(), nullptr, {});
+  out.copy_shared_buffer(
+      x,
+      std::move(strides),
+      {contiguous, row_contiguous, col_contiguous},
+      x.data_size());
+  return out;
+}
+
 } // namespace mlx::core

mlx/backend/common/utils.h

@@ -196,6 +196,9 @@ void shared_buffer_reshape(
     const Strides& out_strides,
     array& out);
 
+// Like the swapaxes op but safe to call in eval_gpu.
+array swapaxes_in_eval(const array& x, int axis1, int axis2);
+
 template <typename T>
 inline std::vector<T> remove_index(std::vector<T> vec, size_t index) {
   vec.erase(std::next(vec.begin(), index));

mlx/backend/cuda/CMakeLists.txt

@@ -6,6 +6,7 @@
 target_sources(
   mlx
   PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/allocator.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/arange.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/arg_reduce.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/binary.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/binary_two.cu
@@ -29,7 +30,7 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/layer_norm.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/logsumexp.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/random.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/reduce/all_reduce.cu
@@ -45,7 +46,8 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/ternary.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/unary.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/quantized.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/affine_quantize.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/quantized/quantized.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/worker.cpp)
 
 if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.9.0)
@@ -105,11 +107,11 @@ if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8.0)
     mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--compress-mode=size>")
 endif()
 
-# Compute capability 7 is required for synchronization between CPU/GPU with
-# managed memory. TODO: Add more architectures for potential performance gain.
-set(MLX_CUDA_ARCHITECTURES
-    "70;80"
-    CACHE STRING "CUDA architectures")
+# Compute capability >= 7.0 is required for synchronization between CPU/GPU with
+# managed memory.
+if(NOT DEFINED MLX_CUDA_ARCHITECTURES)
+  set(MLX_CUDA_ARCHITECTURES "native")
+endif()
 message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}")
 set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES
                                      "${MLX_CUDA_ARCHITECTURES}")

mlx/backend/cuda/arange.cu

@@ -0,0 +1,55 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/device/fp16_math.cuh"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/dtype_utils.h"
+#include "mlx/primitives.h"
+
+#include <nvtx3/nvtx3.hpp>
+#include <thrust/device_ptr.h>
+#include <thrust/transform.h>
+
+namespace mlx::core {
+
+namespace cu {
+
+template <typename T>
+struct Arange {
+  const T start;
+  const T step;
+
+  __device__ T operator()(uint32_t i) const {
+    return start + i * step;
+  }
+};
+
+} // namespace cu
+
+void Arange::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("Arange::eval_gpu");
+  if (out.size() == 0) {
+    return;
+  }
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& encoder = cu::get_command_encoder(stream());
+  encoder.set_output_array(out);
+
+  auto capture = encoder.capture_context();
+  dispatch_int_float_types(out.dtype(), "Arange", [&](auto type_tag) {
+    using CTYPE = MLX_GET_TYPE(type_tag);
+    using OutType = cuda_type_t<CTYPE>;
+    CTYPE step =
+        static_cast<CTYPE>(start_ + step_) - static_cast<CTYPE>(start_);
+    thrust::transform(
+        cu::thrust_policy(encoder.stream()),
+        thrust::counting_iterator<uint32_t>(0),
+        thrust::counting_iterator<uint32_t>(out.data_size()),
+        thrust::device_pointer_cast(out.data<OutType>()),
+        cu::Arange<OutType>{
+            static_cast<OutType>(start_), static_cast<OutType>(step)});
+  });
+}
+
+} // namespace mlx::core

mlx/backend/cuda/arg_reduce.cu

@@ -44,8 +44,11 @@ struct ArgMin {
   }
 
   template <int N>
-  __device__ IndexValPair<T>
-  reduce_many(IndexValPair<T> best, T (&vals)[N], uint32_t offset) {
+  __device__ IndexValPair<T> reduce_many(
+      IndexValPair<T> best,
+      const AlignedVector<T, N>& vals,
+      uint32_t offset) {
+#pragma unroll
     for (int i = 0; i < N; i++) {
       if (vals[i] < best.val) {
         best.val = vals[i];
@@ -74,8 +77,11 @@ struct ArgMax {
   }
 
   template <int N>
-  __device__ IndexValPair<T>
-  reduce_many(IndexValPair<T> best, T (&vals)[N], uint32_t offset) {
+  __device__ IndexValPair<T> reduce_many(
+      IndexValPair<T> best,
+      const AlignedVector<T, N>& vals,
+      uint32_t offset) {
+#pragma unroll
     for (int i = 0; i < N; i++) {
       if (vals[i] > best.val) {
         best.val = vals[i];
@@ -106,16 +112,15 @@ __global__ void arg_reduce_general(
 
   int64_t in_idx = elem_to_loc(index, shape.data(), in_strides.data(), ndim);
   int64_t out_idx = elem_to_loc(index, shape.data(), out_strides.data(), ndim);
+  in += in_idx;
 
   Op op;
   T init = op.init();
   IndexValPair<T> best{0, init};
 
   for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
-    T vals[N_READS];
     auto tid = r * BLOCK_DIM + block.thread_index().x;
-    cub::LoadDirectBlocked(
-        tid, StridedIterator(in + in_idx, axis_stride), vals, axis_size, init);
+    auto vals = load_vector<N_READS>(in, tid, axis_size, axis_stride, init);
     best = op.reduce_many(best, vals, tid * N_READS);
   }
 

mlx/backend/cuda/binary.cu

@@ -28,7 +28,7 @@ __global__ void binary_ss(const In* a, const In* b, Out* out, IdxT size) {
     AlignedVector<Out, N_READS> out_vec;
 #pragma unroll
     for (int i = 0; i < N_READS; ++i) {
-      out_vec.val[i] = Op{}(a[0], b[0]);
+      out_vec[i] = Op{}(a[0], b[0]);
     }
 
     store_vector<N_READS>(out, index, out_vec);
@@ -49,7 +49,7 @@ __global__ void binary_sv(const In* a, const In* b, Out* out, IdxT size) {
     AlignedVector<Out, N_READS> out_vec;
 #pragma unroll
     for (int i = 0; i < N_READS; ++i) {
-      out_vec.val[i] = Op{}(a[0], b_vec.val[i]);
+      out_vec[i] = Op{}(a[0], b_vec[i]);
     }
 
     store_vector<N_READS>(out, index, out_vec);
@@ -70,7 +70,7 @@ __global__ void binary_vs(const In* a, const In* b, Out* out, IdxT size) {
     AlignedVector<Out, N_READS> out_vec;
 #pragma unroll
     for (int i = 0; i < N_READS; ++i) {
-      out_vec.val[i] = Op{}(a_vec.val[i], b[0]);
+      out_vec[i] = Op{}(a_vec[i], b[0]);
     }
 
     store_vector<N_READS>(out, index, out_vec);
@@ -92,7 +92,7 @@ __global__ void binary_vv(const In* a, const In* b, Out* out, IdxT size) {
     AlignedVector<Out, N_READS> out_vec;
 #pragma unroll
     for (int i = 0; i < N_READS; ++i) {
-      out_vec.val[i] = Op{}(a_vec.val[i], b_vec.val[i]);
+      out_vec[i] = Op{}(a_vec[i], b_vec[i]);
     }
 
     store_vector<N_READS>(out, index, out_vec);
@@ -211,12 +211,15 @@ void binary_op_gpu_inplace(
                 int ndim = shape.size();
                 if (ndim <= 3) {
                   dispatch_1_2_3(ndim, [&](auto dims_constant) {
-                    auto kernel = cu::
-                        binary_g_nd<Op, InType, OutType, IdxT, dims_constant()>;
                     auto [num_blocks, block_dims] =
-                        get_launch_args(kernel, out, large());
+                        get_launch_args(out, large());
                     encoder.add_kernel_node(
-                        kernel,
+                        cu::binary_g_nd<
+                            Op,
+                            InType,
+                            OutType,
+                            IdxT,
+                            dims_constant()>,
                         num_blocks,
                         block_dims,
                         a.data<InType>(),
@@ -228,11 +231,9 @@ void binary_op_gpu_inplace(
                         const_param<dims_constant()>(b_strides));
                   });
                 } else {
-                  auto kernel = cu::binary_g<Op, InType, OutType, IdxT>;
-                  auto [num_blocks, block_dims] =
-                      get_launch_args(kernel, out, large());
+                  auto [num_blocks, block_dims] = get_launch_args(out, large());
                   encoder.add_kernel_node(
-                      kernel,
+                      cu::binary_g<Op, InType, OutType, IdxT>,
                       num_blocks,
                       block_dims,
                       a.data<InType>(),
@@ -248,8 +249,7 @@ void binary_op_gpu_inplace(
         } else {
           dispatch_bool(out.data_size() > UINT32_MAX, [&](auto large) {
             using IdxT = std::conditional_t<large(), int64_t, uint32_t>;
-            // TODO: Choose optimized value based on type size.
-            constexpr int N_READS = 4;
+            constexpr int N_READS = 16 / sizeof(InType);
             auto kernel = cu::binary_ss<Op, InType, OutType, IdxT, N_READS>;
             if (bopt == BinaryOpType::ScalarVector) {
               kernel = cu::binary_sv<Op, InType, OutType, IdxT, N_READS>;
@@ -259,12 +259,7 @@ void binary_op_gpu_inplace(
               kernel = cu::binary_vv<Op, InType, OutType, IdxT, N_READS>;
             }
             auto [num_blocks, block_dims] = get_launch_args(
-                kernel,
-                out.data_size(),
-                out.shape(),
-                out.strides(),
-                large(),
-                N_READS);
+                out.data_size(), out.shape(), out.strides(), large(), N_READS);
             encoder.add_kernel_node(
                 kernel,
                 num_blocks,

mlx/backend/cuda/binary_two.cu

@@ -33,8 +33,8 @@ binary_two_ss(const In* a, const In* b, Out* out_a, Out* out_b, IdxT size) {
 #pragma unroll
     for (int i = 0; i < N_READS; ++i) {
       auto out = Op{}(a[0], b[0]);
-      out_a_vec.val[i] = out[0];
-      out_b_vec.val[i] = out[1];
+      out_a_vec[i] = out[0];
+      out_b_vec[i] = out[1];
     }
 
     store_vector<N_READS>(out_a, index, out_a_vec);
@@ -60,9 +60,9 @@ binary_two_sv(const In* a, const In* b, Out* out_a, Out* out_b, IdxT size) {
     AlignedVector<Out, N_READS> out_b_vec;
 #pragma unroll
     for (int i = 0; i < N_READS; ++i) {
-      auto out = Op{}(a[0], b_vec.val[i]);
-      out_a_vec.val[i] = out[0];
-      out_b_vec.val[i] = out[1];
+      auto out = Op{}(a[0], b_vec[i]);
+      out_a_vec[i] = out[0];
+      out_b_vec[i] = out[1];
     }
 
     store_vector<N_READS>(out_a, index, out_a_vec);
@@ -88,9 +88,9 @@ binary_two_vs(const In* a, const In* b, Out* out_a, Out* out_b, IdxT size) {
     AlignedVector<Out, N_READS> out_b_vec;
 #pragma unroll
     for (int i = 0; i < N_READS; ++i) {
-      auto out = Op{}(a_vec.val[i], b[0]);
-      out_a_vec.val[i] = out[0];
-      out_b_vec.val[i] = out[1];
+      auto out = Op{}(a_vec[i], b[0]);
+      out_a_vec[i] = out[0];
+      out_b_vec[i] = out[1];
     }
 
     store_vector<N_READS>(out_a, index, out_a_vec);
@@ -117,9 +117,9 @@ binary_two_vv(const In* a, const In* b, Out* out_a, Out* out_b, IdxT size) {
     AlignedVector<Out, N_READS> out_b_vec;
 #pragma unroll
     for (int i = 0; i < N_READS; ++i) {
-      auto out = Op{}(a_vec.val[i], b_vec.val[i]);
-      out_a_vec.val[i] = out[0];
-      out_b_vec.val[i] = out[1];
+      auto out = Op{}(a_vec[i], b_vec[i]);
+      out_a_vec[i] = out[0];
+      out_b_vec[i] = out[1];
     }
 
     store_vector<N_READS>(out_a, index, out_a_vec);
@@ -227,16 +227,15 @@ void binary_two_op_gpu_inplace(
                 int ndim = shape.size();
                 if (ndim <= 3) {
                   dispatch_1_2_3(ndim, [&](auto dims_constant) {
-                    auto kernel = cu::binary_two_g_nd<
-                        Op,
-                        InType,
-                        OutType,
-                        IdxT,
-                        dims_constant()>;
                     auto [num_blocks, block_dims] =
-                        get_launch_args(kernel, out_a, large());
+                        get_launch_args(out_a, large());
                     encoder.add_kernel_node(
-                        kernel,
+                        cu::binary_two_g_nd<
+                            Op,
+                            InType,
+                            OutType,
+                            IdxT,
+                            dims_constant()>,
                         num_blocks,
                         block_dims,
                         a.data<InType>(),
@@ -249,11 +248,10 @@ void binary_two_op_gpu_inplace(
                         const_param<dims_constant()>(b_strides));
                   });
                 } else {
-                  auto kernel = cu::binary_two_g<Op, InType, OutType, IdxT>;
                   auto [num_blocks, block_dims] =
-                      get_launch_args(kernel, out_a, large());
+                      get_launch_args(out_a, large());
                   encoder.add_kernel_node(
-                      kernel,
+                      cu::binary_two_g<Op, InType, OutType, IdxT>,
                       num_blocks,
                       block_dims,
                       a.data<InType>(),
@@ -270,8 +268,7 @@ void binary_two_op_gpu_inplace(
         } else {
           dispatch_bool(out_a.data_size() > UINT32_MAX, [&](auto large) {
             using IdxT = std::conditional_t<large(), int64_t, uint32_t>;
-            // TODO: Choose optimized value based on type size.
-            constexpr int N_READS = 4;
+            constexpr int N_READS = 16 / sizeof(InType);
             auto kernel = cu::binary_two_ss<Op, InType, OutType, IdxT, N_READS>;
             if (bopt == BinaryOpType::ScalarVector) {
               kernel = cu::binary_two_sv<Op, InType, OutType, IdxT, N_READS>;
@@ -281,7 +278,6 @@ void binary_two_op_gpu_inplace(
               kernel = cu::binary_two_vv<Op, InType, OutType, IdxT, N_READS>;
             }
             auto [num_blocks, block_dims] = get_launch_args(
-                kernel,
                 out_a.data_size(),
                 out_a.shape(),
                 out_a.strides(),

mlx/backend/cuda/compiled.cpp

@@ -104,10 +104,41 @@ struct FusedKernelBuilder {
           "  }\n";
     }
 
+    // Vectorized read loop
+    if (contiguous) {
+      for (size_t i = 0; i < inputs.size(); ++i) {
+        const auto& x = inputs[i];
+        if (is_scalar(x) || is_constant(i)) {
+          continue;
+        }
+        const std::string& xname = namer.get_name(x);
+        std::string type = dtype_to_cuda_type(x.dtype());
+        os += fmt::format(
+            "  auto vec_{0} = load_vector<work_per_thread, {1}>({0} + index, 0, size - index, 0);\n",
+            xname,
+            type);
+      }
+    }
+
+    // Create some space for the outputs
+    for (const auto& x : outputs) {
+      const std::string& xname = namer.get_name(x);
+      std::string type = dtype_to_cuda_type(x.dtype());
+      os += fmt::format(
+          "  AlignedVector<{}, work_per_thread> vec_{};\n", type, xname);
+    }
+
     // Work loop
-    os +=
-        "\n"
-        "  for (int i = 0; i < work_per_thread && index < size; i++) {\n";
+    if (!contiguous) {
+      os +=
+          "\n"
+          "  for (int i = 0; i < work_per_thread && index < size; i++) {\n";
+    } else {
+      os +=
+          "\n"
+          "  #pragma unroll\n"
+          "  for (int i = 0; i < work_per_thread; i++) {\n";
+    }
 
     // Read inputs.
     for (size_t i = 0; i < inputs.size(); ++i) {
@@ -122,7 +153,7 @@ struct FusedKernelBuilder {
       } else if (is_scalar(x)) {
         value = fmt::format("{}[0]", xname);
       } else if (contiguous) {
-        value = fmt::format("{}[index]", xname);
+        value = fmt::format("vec_{}[i]", xname);
       } else {
         value = fmt::format("{}[{}_idx]", xname, xname);
       }
@@ -150,25 +181,30 @@ struct FusedKernelBuilder {
 
     // Write output.
     for (const auto& x : outputs) {
-      os += fmt::format("    {0}[index] = tmp_{0};\n", namer.get_name(x));
+      os += fmt::format("    vec_{0}[i] = tmp_{0};\n", namer.get_name(x));
     }
 
     // End of work loop
-    os +=
-        "\n"
-        "    index++;\n";
     if (!contiguous) {
+      os += "\n";
       for (size_t i = 0; i < inputs.size(); ++i) {
         const auto& x = inputs[i];
         const std::string& xname = namer.get_name(x);
         if (is_scalar(x) || is_constant(i)) {
           continue;
         }
-        os += "    " + xname + "_idx += " + xname + "_strides[NDIM - 1];\n";
+        os += fmt::format("    {0}_idx += {0}_strides[NDIM - 1];\n", xname);
       }
     }
     os += "  }\n";
 
+    // Store the output to global memory
+    for (const auto& x : outputs) {
+      os += fmt::format(
+          "  store_vector({0} + index, 0, vec_{0}, size - index);\n",
+          namer.get_name(x));
+    }
+
     os += "}\n";
   }
 };
@@ -192,6 +228,15 @@ void Compiled::eval_gpu(
   nvtx3::scoped_range r("Compiled::eval_gpu");
   auto& s = stream();
 
+  // Determine the work per thread for the vectorized reads/writes. We take it
+  // as 16 over the max itemsize for the outputs. Another heuristic could be
+  // over the max itemsize of all arrays.
+  int max_size = 1;
+  for (const auto& x : outputs) {
+    max_size = (max_size > x.itemsize()) ? max_size : x.itemsize();
+  }
+  int work_per_thread = 16 / max_size;
+
   cu::JitModule& mod = cu::get_jit_module(s.device, lib_name(), [&]() {
     // Build source code.
     cu::FusedKernelBuilder builder{
@@ -205,28 +250,23 @@ void Compiled::eval_gpu(
     builder.os += "\n} // namespace mlx::core::cu\n";
     // Build kernel names.
     std::vector<std::string> kernel_names;
-    for (auto work_per_thread : std::array<int, 2>{1, 4}) {
-      kernel_names.push_back(fmt::format(
-          "mlx::core::cu::{}_contiguous<uint32_t, {}>",
-          lib_name(),
-          work_per_thread));
-      kernel_names.push_back(fmt::format(
-          "mlx::core::cu::{}_contiguous<int64_t, {}>",
-          lib_name(),
-          work_per_thread));
+    kernel_names.push_back(fmt::format(
+        "mlx::core::cu::{}_contiguous<uint32_t, {}>",
+        lib_name(),
+        work_per_thread));
+    kernel_names.push_back(fmt::format(
+        "mlx::core::cu::{}_contiguous<int64_t, {}>",
+        lib_name(),
+        work_per_thread));
+    for (auto wpt : std::array<int, 2>{1, work_per_thread}) {
       for (int i = 1; i <= MAX_NDIM; ++i) {
         kernel_names.push_back(fmt::format(
-            "mlx::core::cu::{}_strided<{}, uint32_t, {}>",
-            lib_name(),
-            i,
-            work_per_thread));
+            "mlx::core::cu::{}_strided<{}, uint32_t, {}>", lib_name(), i, wpt));
         kernel_names.push_back(fmt::format(
-            "mlx::core::cu::{}_strided<{}, int64_t, {}>",
-            lib_name(),
-            i,
-            work_per_thread));
+            "mlx::core::cu::{}_strided<{}, int64_t, {}>", lib_name(), i, wpt));
       }
     }
+
     return std::make_pair(std::move(builder.os), std::move(kernel_names));
   });
 
@@ -269,7 +309,6 @@ void Compiled::eval_gpu(
   }
 
   // Choose work per thread
-  int work_per_thread = 4;
   if (!contiguous && shape.back() % work_per_thread != 0) {
     work_per_thread = 1;
   }
@@ -294,7 +333,7 @@ void Compiled::eval_gpu(
 
   auto kernel = mod.get_kernel(kernel_name);
   auto [num_blocks, block_dims] =
-      get_launch_args(kernel, outputs[0], large, work_per_thread);
+      get_launch_args(outputs[0], large, work_per_thread);
   encoder.add_kernel_node(kernel, num_blocks, block_dims, args.args());
 }
 

mlx/backend/cuda/conv.cpp

@@ -16,7 +16,6 @@
 #include <nvtx3/nvtx3.hpp>
 
 #include <cassert>
-#include <numeric>
 
 namespace mlx::core {
 
@@ -25,25 +24,34 @@ namespace {
 // Not all engines support it so can not use this API now.
 #define MLX_USE_CUDNN_NATIVE_CUDA_GRAPH_API 0
 
+// Alias for better readability.
+#define CONV_FORWARD CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR
+#define CONV_BACKWARD_INPUT \
+  CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR
+#define CONV_BACKWARD_WEIGHT \
+  CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR
+
 struct ConvCacheKey {
   int device_id;
-  cudnnBackendDescriptorType_t backend_type;
-  cudnnDataType_t cudnn_type;
+  cudnnDataType_t cudnn_dtype;
   std::array<int, MAX_NDIM> input_shape;
-  std::array<int, MAX_NDIM> filter_shape;
+  std::array<int, MAX_NDIM> weight_shape;
+  std::array<int, MAX_NDIM> stride;
   std::array<int, MAX_NDIM> padding_lo;
   std::array<int, MAX_NDIM> padding_hi;
-  std::array<int, MAX_NDIM> stride;
   std::array<int, MAX_NDIM> dilation;
   int groups;
+  bool flip;
   uint8_t input_alignment;
-  uint8_t filter_alignment;
+  uint8_t weight_alignment;
   uint8_t output_alignment;
 };
 
 auto& conv_cache() {
-  static LRUBytesKeyCache<ConvCacheKey, cudnn_frontend::ExecutionPlan> cache(
-      /* capacity */ 128);
+  static LRUBytesKeyCache<
+      ConvCacheKey,
+      std::pair<cudnnBackendDescriptorType_t, cudnn_frontend::ExecutionPlan>>
+      cache(/* capacity */ 128);
   return cache;
 }
 
@@ -162,17 +170,17 @@ cudnn_frontend::EngineConfigList get_engine_configs(
 bool execute_plan(
     cu::CommandEncoder& encoder,
     cudnn_frontend::ExecutionPlan& plan,
-    const array& in,
-    const array& wt,
-    array& out) {
+    array& x,
+    array& w,
+    array& y) {
   int workspace_size = plan.getWorkspaceSize();
   array workspace(allocator::malloc(workspace_size), {workspace_size}, uint8);
 
   int64_t uids[3] = {'x', 'w', 'y'};
   void* data_ptrs[3] = {
-      const_cast<void*>(in.data<void>()),
-      const_cast<void*>(wt.data<void>()),
-      out.data<void>(),
+      x.data<void>(),
+      w.data<void>(),
+      y.data<void>(),
   };
 
   auto variantPack = cudnn_frontend::VariantPackBuilder()
@@ -212,46 +220,154 @@ bool execute_plan(
 
 bool try_engines(
     cu::CommandEncoder& encoder,
-    cudnn_frontend::EngineConfigList& configs,
     const ConvCacheKey& cache_key,
+    cudnnBackendDescriptorType_t backend_type,
+    cudnn_frontend::EngineConfigList& configs,
     const std::string& op_graph_tag,
-    const array& in,
-    const array& wt,
-    array& out) {
+    array& x,
+    array& w,
+    array& y) {
   for (auto& config : configs) {
     try {
       auto plan = cudnn_frontend::ExecutionPlanBuilder()
                       .setHandle(encoder.device().cudnn_handle())
                       .setEngineConfig(config, op_graph_tag)
                       .build();
-      if (execute_plan(encoder, plan, in, wt, out)) {
-        conv_cache().emplace(cache_key, std::move(plan));
+      if (execute_plan(encoder, plan, x, w, y)) {
+        conv_cache().emplace(
+            cache_key, std::make_pair(backend_type, std::move(plan)));
         return true;
       }
-    } catch (cudnn_frontend::cudnnException&) {
+    } catch (cudnn_frontend::cudnnException& error) {
+      if (error.getCudnnStatus() != CUDNN_STATUS_NOT_SUPPORTED) {
+        throw;
+      }
     }
   }
   return false;
 }
 
-} // namespace
+auto get_conv_op_settings(
+    cudnnBackendDescriptorType_t backend_type,
+    array& x,
+    array& w,
+    array& y,
+    const std::vector<int>& kernel_strides,
+    const std::vector<int>& padding_lo_,
+    const std::vector<int>& padding_hi_,
+    const std::vector<int>& kernel_dilation,
+    const std::vector<int>& input_dilation) {
+  auto padding_lo = convert_vector<int64_t>(padding_lo_);
+  auto padding_hi = convert_vector<int64_t>(padding_hi_);
 
-void Convolution::eval_gpu(const std::vector<array>& inputs, array& out) {
-  nvtx3::scoped_range r("Convolution::eval_gpu");
-  if (out.size() == 0) {
-    return;
+  if (backend_type == CONV_BACKWARD_INPUT) {
+    for (int i = 0; i < padding_lo.size(); ++i) {
+      int wt_size = 1 + kernel_dilation[i] * (w.shape(1 + i) - 1);
+      padding_lo[i] = wt_size - padding_lo[i] - 1;
+      int in_size = 1 + kernel_strides[i] * (x.shape(1 + i) - 1);
+      int out_size = 1 + input_dilation[i] * (y.shape(1 + i) - 1);
+      padding_hi[i] = out_size - in_size + padding_hi[i];
+    }
+    return std::make_tuple(
+        convert_vector<int64_t>(input_dilation),
+        std::move(padding_lo),
+        std::move(padding_hi),
+        convert_vector<int64_t>(kernel_dilation));
+
+  } else if (backend_type == CONV_BACKWARD_WEIGHT) {
+    padding_hi = padding_lo;
+    return std::make_tuple(
+        convert_vector<int64_t>(kernel_dilation),
+        std::move(padding_lo),
+        std::move(padding_hi),
+        convert_vector<int64_t>(kernel_strides));
+
+  } else {
+    return std::make_tuple(
+        convert_vector<int64_t>(kernel_strides),
+        std::move(padding_lo),
+        std::move(padding_hi),
+        convert_vector<int64_t>(kernel_dilation));
+  }
+}
+
+std::optional<cudnn_frontend::OperationGraph> build_op_graph(
+    cu::CommandEncoder& encoder,
+    cudnnBackendDescriptorType_t backend_type,
+    Dtype dtype,
+    array& x,
+    array& w,
+    array& y,
+    const std::vector<int64_t>& stride,
+    const std::vector<int64_t>& padding_lo,
+    const std::vector<int64_t>& padding_hi,
+    const std::vector<int64_t>& dilation) {
+  try {
+    auto compute_dtype = (dtype == float16 || dtype == bfloat16)
+        ? CUDNN_DATA_FLOAT
+        : dtype_to_cudnn_type(dtype);
+    auto conv_desc = cudnn_frontend::ConvDescBuilder()
+                         .setDataType(compute_dtype)
+                         .setMathMode(CUDNN_CROSS_CORRELATION)
+                         .setNDims(stride.size())
+                         .setStrides(stride.size(), stride.data())
+                         .setPrePadding(padding_lo.size(), padding_lo.data())
+                         .setPostPadding(padding_hi.size(), padding_hi.data())
+                         .setDilation(dilation.size(), dilation.data())
+                         .build();
+
+    auto op = cudnn_frontend::OperationBuilder(backend_type)
+                  .setxDesc(build_tensor('x', x))
+                  .setwDesc(build_tensor('w', w))
+                  .setyDesc(build_tensor('y', y))
+                  .setcDesc(conv_desc)
+                  .build();
+
+    std::array<cudnn_frontend::Operation const*, 1> ops = {&op};
+    return cudnn_frontend::OperationGraphBuilder()
+        .setHandle(encoder.device().cudnn_handle())
+        .setOperationGraph(ops.size(), ops.data())
+        .build();
+  } catch (cudnn_frontend::cudnnException& error) {
+    if (error.getCudnnStatus() != CUDNN_STATUS_BAD_PARAM) {
+      throw;
+    }
+    return std::nullopt;
+  }
+}
+
+// Do necessary transposes and copies to prepare the inputs and outputs for
+// building the cuDNN conv op. It is safe to be called multiple times in one
+// eval_gpu, with cost of possible redundant copies.
+std::tuple<array, array, array> prepare_args(
+    cu::CommandEncoder& encoder,
+    cudnnBackendDescriptorType_t backend_type,
+    array in,
+    array wt,
+    array out,
+    Stream s) {
+  // Transpose the args depending on the backend type.
+  // TODO: Handle groups.
+  if (backend_type == CONV_BACKWARD_INPUT) {
+    wt = swapaxes_in_eval(wt, 0, -1);
+  } else if (backend_type == CONV_BACKWARD_WEIGHT) {
+    in = swapaxes_in_eval(in, 0, -1);
+    wt = swapaxes_in_eval(wt, 0, -1);
+    // Create a contiguous array that shares the data with |out|, but with dim
+    // C_in and C_out swapped.
+    Shape shape(out.shape());
+    std::swap(shape.front(), shape.back());
+    Strides strides(shape.size(), 1);
+    for (int i = shape.size() - 2; i >= 0; --i) {
+      strides[i] = shape[i + 1] * strides[i + 1];
+    }
+    array intermediate(std::move(shape), out.dtype(), nullptr, {});
+    intermediate.copy_shared_buffer(
+        out, std::move(strides), {true, true, false}, out.data_size());
+    out = intermediate;
   }
 
-  assert(inputs.size() == 2);
-  array in = inputs[0];
-  array wt = inputs[1];
-  out.set_data(allocator::malloc(out.nbytes()));
-
-  auto& s = stream();
-  auto& encoder = cu::get_command_encoder(s);
-
   // cuDNN requires contiguous input.
-  // TODO: Handle NCHW format specially.
   if (!in.flags().row_contiguous) {
     in = contiguous_copy_gpu(in, s);
     encoder.add_temporary(in);
@@ -261,80 +377,170 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out) {
     encoder.add_temporary(wt);
   }
 
+  return {std::move(in), std::move(wt), std::move(out)};
+}
+
+// Get the x/w/y args from the in/wt/out args depending on backend type.
+inline std::tuple<array&, array&, array&> dispatch_args(
+    cudnnBackendDescriptorType_t backend_type,
+    array& in,
+    array& wt,
+    array& out) {
+  switch (backend_type) {
+    case CONV_BACKWARD_INPUT:
+      return {out, wt, in};
+    case CONV_BACKWARD_WEIGHT:
+      return {in, out, wt};
+    default:
+      return {in, wt, out};
+  }
+}
+
+// Register inputs and outputs before actually running conv op. Can only be
+// called once per eval_gpu.
+void register_args(
+    cu::CommandEncoder& encoder,
+    cudnnBackendDescriptorType_t backend_type,
+    array& in,
+    array& wt,
+    array& intermediate_out,
+    array& final_out) {
   encoder.set_input_array(in);
   encoder.set_input_array(wt);
-  encoder.set_output_array(out);
+  encoder.set_output_array(final_out);
 
-  auto backend_type = CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR;
-  auto cudnn_type = dtype_to_cudnn_type(in.dtype());
+  if (backend_type == CONV_BACKWARD_WEIGHT) {
+    // Turn |out| into a strided array, which will have C_in and C_out swapped
+    // in vjp and the final |grad_weight| will then be contiguous.
+    Strides strides = intermediate_out.strides();
+    std::swap(strides.front(), strides.back());
+    final_out.copy_shared_buffer(
+        intermediate_out,
+        std::move(strides),
+        {false, false, false},
+        intermediate_out.data_size());
+  }
+}
+
+} // namespace
+
+void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
+  nvtx3::scoped_range r("Convolution::eval_gpu");
+  if (out_.size() == 0) {
+    return;
+  }
+
+  assert(inputs.size() == 2);
+  array in = inputs[0];
+  array wt = inputs[1];
+  array out = out_;
+  out.set_data(allocator::malloc(out.nbytes()));
+  Dtype dtype = out.dtype();
+
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
 
   // Search cache.
   ConvCacheKey cache_key{
       encoder.device().cuda_device(),
-      backend_type,
-      cudnn_type,
+      dtype_to_cudnn_type(dtype),
       fixed_vector(in.shape()),
       fixed_vector(wt.shape()),
+      fixed_vector(kernel_strides_),
       fixed_vector(padding_lo_),
       fixed_vector(padding_hi_),
-      fixed_vector(kernel_strides_),
       fixed_vector(kernel_dilation_),
       groups_,
+      flip_,
       get_alignment(in),
       get_alignment(wt),
       get_alignment(out)};
   if (auto it = conv_cache().find(cache_key); it != conv_cache().end()) {
-    if (!execute_plan(encoder, it->second, in, wt, out)) {
-      throw std::runtime_error("Cached convolution plan failed to execute.");
+    auto& [backend_type, plan] = it->second;
+    std::tie(in, wt, out) = prepare_args(encoder, backend_type, in, wt, out, s);
+    register_args(encoder, backend_type, in, wt, out, out_);
+    auto [x, w, y] = dispatch_args(backend_type, in, wt, out);
+    if (!execute_plan(encoder, plan, x, w, y)) {
+      throw std::runtime_error("[conv] Cached plan failed to execute.");
     }
     return;
   }
 
-  // Build operation graph.
-  auto compute_data_type = (in.dtype() == float16 || in.dtype() == bfloat16)
-      ? CUDNN_DATA_FLOAT
-      : cudnn_type;
+  // There is no reliable way to deduce the proper cuDNN backend for the
+  // convolution, so we make a best guess and then try.
+  std::vector<cudnnBackendDescriptorType_t> try_backends;
+  if (flip_) {
+    // When weight is flipped, we assume it is backward input convolution.
+    try_backends.push_back(CONV_BACKWARD_INPUT);
+  } else {
+    // Otherwise it could be backward weight convolution or forward convolution,
+    // mathematically there is no difference so we have to use heuristics.
+    // Empirically backward convolutions have large kernel dimensions, and
+    // usually have |in| and |wt| transposed.
+    if (!in.flags().row_contiguous && !wt.flags().row_contiguous &&
+        wt.shape(2) > out.shape(2)) {
+      try_backends = {CONV_BACKWARD_WEIGHT, CONV_FORWARD};
+    } else {
+      try_backends = {CONV_FORWARD, CONV_BACKWARD_WEIGHT};
+    }
+  }
 
-  auto stride = convert_vector<int64_t>(kernel_strides_);
-  auto padding_lo = convert_vector<int64_t>(padding_lo_);
-  auto padding_hi = convert_vector<int64_t>(padding_hi_);
-  auto dilation = convert_vector<int64_t>(kernel_dilation_);
+  // Try to build op graph.
+  cudnnBackendDescriptorType_t backend_type;
+  std::optional<cudnn_frontend::OperationGraph> op_graph;
+  for (auto try_backend : try_backends) {
+    auto [in_copy, wt_copy, out_copy] =
+        prepare_args(encoder, try_backend, in, wt, out, s);
+    auto [x, w, y] = dispatch_args(try_backend, in_copy, wt_copy, out_copy);
+    auto [stride, padding_lo, padding_hi, dilation] = get_conv_op_settings(
+        try_backend,
+        x,
+        w,
+        y,
+        kernel_strides_,
+        padding_lo_,
+        padding_hi_,
+        kernel_dilation_,
+        input_dilation_);
+    op_graph = build_op_graph(
+        encoder,
+        try_backend,
+        dtype,
+        x,
+        w,
+        y,
+        stride,
+        padding_lo,
+        padding_hi,
+        dilation);
+    if (op_graph) {
+      backend_type = try_backend;
+      in = std::move(in_copy);
+      wt = std::move(wt_copy);
+      out = std::move(out_copy);
+      break;
+    }
+  }
+  if (!op_graph) {
+    throw std::runtime_error("[conv] Can not build op graph.");
+  }
 
-  auto conv_desc = cudnn_frontend::ConvDescBuilder()
-                       .setDataType(compute_data_type)
-                       .setMathMode(CUDNN_CROSS_CORRELATION)
-                       .setNDims(stride.size())
-                       .setStrides(stride.size(), stride.data())
-                       .setPrePadding(padding_lo.size(), padding_lo.data())
-                       .setPostPadding(padding_hi.size(), padding_hi.data())
-                       .setDilation(dilation.size(), dilation.data())
-                       .build();
-
-  auto op = cudnn_frontend::OperationBuilder(backend_type)
-                .setxDesc(build_tensor('x', in))
-                .setwDesc(build_tensor('w', wt))
-                .setyDesc(build_tensor('y', out))
-                .setcDesc(conv_desc)
-                .build();
-
-  std::array<cudnn_frontend::Operation const*, 1> ops = {&op};
-  auto op_graph = cudnn_frontend::OperationGraphBuilder()
-                      .setHandle(encoder.device().cudnn_handle())
-                      .setOperationGraph(ops.size(), ops.data())
-                      .build();
+  // Get ready to execute the graph.
+  register_args(encoder, backend_type, in, wt, out, out_);
 
   // Try to run plans based on heuristics.
-  auto configs = get_engine_configs(backend_type, in.dtype(), op_graph);
-  auto op_graph_tag = op_graph.getTag();
-  if (try_engines(encoder, configs, cache_key, op_graph_tag, in, wt, out)) {
+  auto configs = get_engine_configs(backend_type, dtype, *op_graph);
+  auto tag = op_graph->getTag();
+  auto [x, w, y] = dispatch_args(backend_type, in, wt, out);
+  if (try_engines(encoder, cache_key, backend_type, configs, tag, x, w, y)) {
     return;
   }
   // Then try fallback plans.
-  configs = get_engine_configs(backend_type, in.dtype(), op_graph);
-  if (try_engines(encoder, configs, cache_key, op_graph_tag, in, wt, out)) {
+  configs = get_engine_configs(backend_type, dtype, *op_graph);
+  if (try_engines(encoder, cache_key, backend_type, configs, tag, x, w, y)) {
     return;
   }
-  throw std::runtime_error("Unable to find an engine for convolution.");
+  throw std::runtime_error("[conv] Unable to find a working engine.");
 }
 
 } // namespace mlx::core

mlx/backend/cuda/copy/copy_contiguous.cu

@@ -22,7 +22,7 @@ __global__ void copy_s(const In* in, Out* out, IdxT size) {
     AlignedVector<Out, N_READS> out_vec;
 #pragma unroll
     for (int i = 0; i < N_READS; ++i) {
-      out_vec.val[i] = cast_to<Out>(in[0]);
+      out_vec[i] = cast_to<Out>(in[0]);
     }
 
     store_vector<N_READS>(out, index, out_vec);
@@ -43,7 +43,7 @@ __global__ void copy_v(const In* in, Out* out, IdxT size) {
     AlignedVector<Out, N_READS> out_vec;
 #pragma unroll
     for (int i = 0; i < N_READS; ++i) {
-      out_vec.val[i] = cast_to<Out>(in_vec.val[i]);
+      out_vec[i] = cast_to<Out>(in_vec[i]);
     }
 
     store_vector<N_READS>(out, index, out_vec);
@@ -65,19 +65,13 @@ void copy_contiguous(
         using InType = cuda_type_t<MLX_GET_TYPE(in_type_tag)>;
         using OutType = cuda_type_t<MLX_GET_TYPE(out_type_tag)>;
         using IdxT = std::conditional_t<large(), int64_t, uint32_t>;
-        // TODO: Choose optimized value based on type size.
-        constexpr int N_READS = 4;
+        constexpr int N_READS = 16 / sizeof(InType);
         auto kernel = cu::copy_s<InType, OutType, IdxT, N_READS>;
         if (ctype == CopyType::Vector) {
           kernel = cu::copy_v<InType, OutType, IdxT, N_READS>;
         }
         auto [num_blocks, block_dims] = get_launch_args(
-            kernel,
-            out.data_size(),
-            out.shape(),
-            out.strides(),
-            large(),
-            N_READS);
+            out.data_size(), out.shape(), out.strides(), large(), N_READS);
         encoder.add_kernel_node(
             kernel,
             num_blocks,

mlx/backend/cuda/copy/copy_general.cu

@@ -71,12 +71,10 @@ void copy_general(
               data_size *= s;
             if (ndim <= 3) {
               dispatch_1_2_3(ndim, [&](auto ndim_constant) {
-                auto kernel =
-                    cu::copy_gg_nd<InType, OutType, IdxT, ndim_constant()>;
-                auto [num_blocks, block_dims] = get_launch_args(
-                    kernel, data_size, shape, out.strides(), large());
+                auto [num_blocks, block_dims] =
+                    get_launch_args(data_size, shape, out.strides(), large());
                 encoder.add_kernel_node(
-                    kernel,
+                    cu::copy_gg_nd<InType, OutType, IdxT, ndim_constant()>,
                     num_blocks,
                     block_dims,
                     in_ptr,
@@ -87,11 +85,10 @@ void copy_general(
                     const_param<ndim_constant()>(strides_out));
               });
             } else { // ndim >= 4
-              auto kernel = cu::copy_gg<InType, OutType, IdxT>;
-              auto [num_blocks, block_dims] = get_launch_args(
-                  kernel, data_size, shape, out.strides(), large());
+              auto [num_blocks, block_dims] =
+                  get_launch_args(data_size, shape, out.strides(), large());
               encoder.add_kernel_node(
-                  kernel,
+                  cu::copy_gg<InType, OutType, IdxT>,
                   num_blocks,
                   block_dims,
                   in_ptr,

mlx/backend/cuda/copy/copy_general_dynamic.cu

@@ -74,12 +74,13 @@ void copy_general_dynamic(
             int ndim = shape.size();
             if (ndim <= 3) {
               dispatch_1_2_3(ndim, [&](auto dims_constant) {
-                auto kernel = cu::
-                    copy_gg_dynamic_nd<InType, OutType, IdxT, dims_constant()>;
-                auto [num_blocks, block_dims] =
-                    get_launch_args(kernel, out, large());
+                auto [num_blocks, block_dims] = get_launch_args(out, large());
                 encoder.add_kernel_node(
-                    kernel,
+                    cu::copy_gg_dynamic_nd<
+                        InType,
+                        OutType,
+                        IdxT,
+                        dims_constant()>,
                     num_blocks,
                     block_dims,
                     in_ptr,
@@ -92,11 +93,9 @@ void copy_general_dynamic(
                     dynamic_offset_out.data<int64_t>());
               });
             } else { // ndim >= 4
-              auto kernel = cu::copy_gg_dynamic<InType, OutType, IdxT>;
-              auto [num_blocks, block_dims] =
-                  get_launch_args(kernel, out, large());
+              auto [num_blocks, block_dims] = get_launch_args(out, large());
               encoder.add_kernel_node(
-                  kernel,
+                  cu::copy_gg_dynamic<InType, OutType, IdxT>,
                   num_blocks,
                   block_dims,
                   in_ptr,

mlx/backend/cuda/copy/copy_general_input.cu

@@ -63,12 +63,9 @@ void copy_general_input(
             int ndim = shape.size();
             if (ndim <= 3) {
               dispatch_1_2_3(ndim, [&](auto dims_constant) {
-                auto kernel =
-                    cu::copy_g_nd<InType, OutType, IdxT, dims_constant()>;
-                auto [num_blocks, block_dims] =
-                    get_launch_args(kernel, out, large());
+                auto [num_blocks, block_dims] = get_launch_args(out, large());
                 encoder.add_kernel_node(
-                    kernel,
+                    cu::copy_g_nd<InType, OutType, IdxT, dims_constant()>,
                     num_blocks,
                     block_dims,
                     in_ptr,
@@ -78,11 +75,9 @@ void copy_general_input(
                     const_param<dims_constant()>(strides_in));
               });
             } else { // ndim >= 4
-              auto kernel = cu::copy_g<InType, OutType, IdxT>;
-              auto [num_blocks, block_dims] =
-                  get_launch_args(kernel, out, large());
+              auto [num_blocks, block_dims] = get_launch_args(out, large());
               encoder.add_kernel_node(
-                  kernel,
+                  cu::copy_g<InType, OutType, IdxT>,
                   num_blocks,
                   block_dims,
                   in_ptr,

mlx/backend/cuda/device.cpp

@@ -1,6 +1,7 @@
 // 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/utils.h"
 
@@ -54,6 +55,10 @@ Device::Device(int device) : device_(device) {
   CHECK_CUBLAS_ERROR(cublasLtCreate(&lt_));
   // The cudnn handle is used by Convolution.
   CHECK_CUDNN_ERROR(cudnnCreate(&cudnn_));
+
+  // Initialize the jit module cache here ensures it is not
+  // unloaded before any evaluation is done
+  get_jit_module_cache();
 }
 
 Device::~Device() {
@@ -92,23 +97,6 @@ CommandEncoder::CaptureContext::~CaptureContext() {
   if (discard) {
     return;
   }
-
-  // Extract and add as single kernel node when possible.
-  size_t num_nodes;
-  CHECK_CUDA_ERROR(cudaGraphGetNodes(graph, NULL, &num_nodes));
-  if (num_nodes == 1) {
-    cudaGraphNode_t captured_node;
-    CHECK_CUDA_ERROR(cudaGraphGetNodes(graph, &captured_node, &num_nodes));
-    cudaGraphNodeType type;
-    CHECK_CUDA_ERROR(cudaGraphNodeGetType(captured_node, &type));
-    if (type == cudaGraphNodeTypeKernel) {
-      CUDA_KERNEL_NODE_PARAMS params;
-      CHECK_CUDA_ERROR(cuGraphKernelNodeGetParams(captured_node, &params));
-      enc.add_kernel_node(params);
-      return;
-    }
-  }
-  // Otherwise add the captured graph as subgraph.
   enc.add_graph_node(graph);
 }
 
@@ -281,6 +269,7 @@ void CommandEncoder::add_graph_node(cudaGraph_t child) {
 }
 
 void CommandEncoder::commit() {
+  nvtx3::scoped_range r("CommandEncoder::commit");
   if (!temporaries_.empty()) {
     add_completed_handler([temporaries = std::move(temporaries_)]() {});
   }
@@ -330,6 +319,7 @@ void CommandEncoder::commit() {
     // Reset state
     node_count_ = 0;
     graph_node_count_ = 0;
+    empty_node_count_ = 0;
     from_nodes_.clear();
     to_nodes_.clear();
     graph_key_.clear();

mlx/backend/cuda/device/arange.cuh

@@ -1,15 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-namespace mlx::core::cu {
-
-template <typename T>
-struct Arange {
-  const T start;
-  const T step;
-
-  __device__ T operator()(uint32_t i) const {
-    return start + i * step;
-  }
-};
-
-} // namespace mlx::core::cu

mlx/backend/cuda/device/atomic_ops.cuh

@@ -49,11 +49,7 @@ inline __device__ void atomic_add(__half* out, __half val) {
 }
 
 inline __device__ void atomic_add(complex64_t* out, complex64_t val) {
-#if __CUDA_ARCH__ < 900
   atomic_add_general(out, val);
-#else
-  atomicAdd(out, val);
-#endif
 }
 
 inline __device__ void atomic_add(__nv_bfloat16* out, __nv_bfloat16 val) {

mlx/backend/cuda/device/utils.cuh

@@ -32,36 +32,119 @@ using Strides = cuda::std::array<int64_t, MAX_NDIM>;
 template <typename T, int N>
 struct alignas(sizeof(T) * N) AlignedVector {
   T val[N];
+
+  __device__ T& operator[](int i) {
+    return val[i];
+  }
+
+  __device__ T operator[](int i) const {
+    return val[i];
+  }
 };
 
 template <int N, typename T>
-inline __device__ AlignedVector<T, N> load_vector(
+inline __host__ __device__ bool is_aligned(T* x) {
+  return (reinterpret_cast<uintptr_t>(x) % (N * sizeof(T))) == 0;
+}
+
+template <int N, typename T>
+inline __device__ AlignedVector<T, N> unsafe_load_vector(
     const T* ptr,
     uint32_t offset) {
   auto* from = reinterpret_cast<const AlignedVector<T, N>*>(ptr);
   return from[offset];
 }
 
+template <int N, typename T>
+inline __device__ AlignedVector<T, N> load_vector(
+    const T* ptr,
+    uint32_t offset) {
+  if (is_aligned<N>(ptr)) {
+    auto* from = reinterpret_cast<const AlignedVector<T, N>*>(ptr);
+    return from[offset];
+  } else {
+    AlignedVector<T, N> v;
+#pragma unroll
+    for (int i = 0; i < N; ++i) {
+      v[i] = ptr[offset * N + i];
+    }
+    return v;
+  }
+}
+
+template <int N, typename T, typename SizeT>
+inline __device__ AlignedVector<T, N>
+load_vector(const T* ptr, uint32_t offset, SizeT size, T fallback) {
+  if (is_aligned<N>(ptr) && (offset + 1) * N <= size) {
+    auto* from = reinterpret_cast<const AlignedVector<T, N>*>(ptr);
+    return from[offset];
+  } else {
+    AlignedVector<T, N> v;
+#pragma unroll
+    for (int i = 0; i < N; ++i) {
+      v[i] = (N * offset + i) < size ? ptr[offset * N + i] : fallback;
+    }
+    return v;
+  }
+}
+
+template <int N, typename T, typename SizeT>
+inline __device__ AlignedVector<T, N> load_vector(
+    const T* ptr,
+    uint32_t offset,
+    SizeT size,
+    int64_t stride,
+    T fallback) {
+  if (is_aligned<N>(ptr) && stride == 1 && (offset + 1) * N <= size) {
+    auto* from = reinterpret_cast<const AlignedVector<T, N>*>(ptr);
+    return from[offset];
+  } else {
+    AlignedVector<T, N> v;
+#pragma unroll
+    for (int i = 0; i < N; ++i) {
+      v[i] =
+          (N * offset + i) < size ? ptr[stride * (offset * N + i)] : fallback;
+    }
+    return v;
+  }
+}
+
 template <int N, typename T>
 inline __device__ void
-store_vector(T* ptr, uint32_t offset, const AlignedVector<T, N>& vec) {
+unsafe_store_vector(T* ptr, uint32_t offset, const AlignedVector<T, N>& vec) {
   auto* to = reinterpret_cast<AlignedVector<T, N>*>(ptr);
   to[offset] = vec;
 }
 
-// Helper for accessing strided data.
-template <typename T>
-struct StridedIterator {
-  T it;
-  int64_t stride;
-
-  __host__ __device__ StridedIterator(T it, int64_t stride)
-      : it(it), stride(stride) {}
-
-  __host__ __device__ auto operator[](int i) const {
-    return it[i * stride];
+template <int N, typename T>
+inline __device__ void
+store_vector(T* ptr, uint32_t offset, const AlignedVector<T, N>& vec) {
+  if (is_aligned<N>(ptr)) {
+    auto* to = reinterpret_cast<AlignedVector<T, N>*>(ptr);
+    to[offset] = vec;
+  } else {
+#pragma unroll
+    for (int i = 0; i < N; ++i) {
+      ptr[offset * N + i] = vec[i];
+    }
   }
-};
+}
+
+template <int N, typename T, typename SizeT>
+inline __device__ void store_vector(
+    T* ptr,
+    uint32_t offset,
+    const AlignedVector<T, N>& vec,
+    SizeT size) {
+  if (is_aligned<N>(ptr) && (offset + 1) * N <= size) {
+    auto* to = reinterpret_cast<AlignedVector<T, N>*>(ptr);
+    to[offset] = vec;
+  } else {
+    for (int i = 0; (offset * N + i) < size && i < N; ++i) {
+      ptr[offset * N + i] = vec[i];
+    }
+  }
+}
 
 ///////////////////////////////////////////////////////////////////////////////
 // Type limits utils

mlx/backend/cuda/eval.cpp

@@ -36,18 +36,15 @@ void eval(array& arr) {
 
   auto& encoder = cu::get_command_encoder(arr.primitive().stream());
   // Keep used buffers alive until kernel finishes running.
-  std::unordered_set<std::shared_ptr<array::Data>> buffers;
   for (auto& in : arr.inputs()) {
-    buffers.insert(in.data_shared_ptr());
+    // Except for the donated one.
+    if (in.data_shared_ptr() != arr.data_shared_ptr()) {
+      encoder.add_temporary(in);
+    }
   }
   for (auto& s : arr.siblings()) {
-    buffers.insert(s.data_shared_ptr());
+    encoder.add_temporary(s);
   }
-  // Remove the output if it was donated to by an input.
-  if (auto it = buffers.find(arr.data_shared_ptr()); it != buffers.end()) {
-    buffers.erase(it);
-  }
-  encoder.add_completed_handler([buffers = std::move(buffers)]() {});
   encoder.maybe_commit();
 }
 

mlx/backend/cuda/gemms/gemv.cu

@@ -11,7 +11,6 @@ namespace mlx::core::cu {
 
 namespace cg = cooperative_groups;
 
-static constexpr int n_per_thread = 4;
 static constexpr int rows_per_block = 8;
 
 template <typename T, int rows_per_block, int n_per_thread>
@@ -28,12 +27,13 @@ gemv_impl(const T* mat, const T* vec, T* out, int rows, int cols) {
     float sum = 0.0f;
     for (int col = n_per_thread * warp.thread_rank(); col < cols;
          col += (WARP_SIZE * n_per_thread)) {
-      auto local_mat = load_vector<n_per_thread>(mat + row * cols + col, 0);
-      auto local_vec = load_vector<n_per_thread>(vec + col, 0);
+      auto local_mat =
+          unsafe_load_vector<n_per_thread>(mat + row * cols + col, 0);
+      auto local_vec = unsafe_load_vector<n_per_thread>(vec + col, 0);
 #pragma unroll
       for (int j = 0; j < n_per_thread; ++j) {
-        sum += static_cast<float>(local_mat.val[j]) *
-            static_cast<float>(local_vec.val[j]);
+        sum +=
+            static_cast<float>(local_mat[j]) * static_cast<float>(local_vec[j]);
       }
     }
 
@@ -74,8 +74,22 @@ __global__ void gemv_batched(
 }
 
 bool can_use_gemv(int M, int N, int K, bool a_transposed, bool b_transposed) {
-  return K % (WARP_SIZE * n_per_thread) == 0 &&
-      ((M == 1 && b_transposed) || (N == 1 && !a_transposed));
+  return K % 32 == 0 && ((M == 1 && b_transposed) || (N == 1 && !a_transposed));
+}
+
+template <typename F>
+void dispatch_n_per_thread(int n_per_thread, F&& f) {
+  switch (n_per_thread) {
+    case 1:
+      f(std::integral_constant<int, 1>{});
+      break;
+    case 2:
+      f(std::integral_constant<int, 2>{});
+      break;
+    case 4:
+      f(std::integral_constant<int, 4>{});
+      break;
+  }
 }
 
 void gemv(
@@ -114,33 +128,43 @@ void gemv(
       rows = M;
     }
     uint32_t num_blocks_x = (rows + rows_per_block - 1) / rows_per_block;
-    if (batch_count == 1) {
-      auto kernel = gemv_single<DataType, rows_per_block, n_per_thread>;
-      encoder.add_kernel_node(
-          kernel,
-          num_blocks_x,
-          block_dims,
-          mat,
-          vec,
-          out.data<DataType>(),
-          rows,
-          cols);
+    int n_per_t;
+    if (K % 128 == 0 && is_aligned<4>(mat) && is_aligned<4>(vec)) {
+      n_per_t = 4;
+    } else if (K % 64 == 0 && is_aligned<2>(mat) && is_aligned<2>(vec)) {
+      n_per_t = 2;
     } else {
-      auto kernel = gemv_batched<DataType, rows_per_block, n_per_thread>;
-      encoder.add_kernel_node(
-          kernel,
-          dim3{num_blocks_x, batch_count},
-          block_dims,
-          mat,
-          vec,
-          out.data<DataType>(),
-          rows,
-          cols,
-          const_param(batch_shape),
-          mat_strides,
-          vec_strides,
-          batch_shape.size());
+      n_per_t = 1;
     }
+    dispatch_n_per_thread(n_per_t, [&](auto n_per_thread) {
+      if (batch_count == 1) {
+        auto kernel = gemv_single<DataType, rows_per_block, n_per_thread()>;
+        encoder.add_kernel_node(
+            kernel,
+            num_blocks_x,
+            block_dims,
+            mat,
+            vec,
+            out.data<DataType>(),
+            rows,
+            cols);
+      } else {
+        auto kernel = gemv_batched<DataType, rows_per_block, n_per_thread()>;
+        encoder.add_kernel_node(
+            kernel,
+            dim3{num_blocks_x, batch_count},
+            block_dims,
+            mat,
+            vec,
+            out.data<DataType>(),
+            rows,
+            cols,
+            const_param(batch_shape),
+            mat_strides,
+            vec_strides,
+            batch_shape.size());
+      }
+    });
   });
 }
 

mlx/backend/cuda/indexing.cpp

@@ -128,7 +128,7 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
   encoder.set_output_array(out);
 
   auto kernel = mod.get_kernel(kernel_name);
-  auto [num_blocks, block_dims] = get_launch_args(kernel, out, large);
+  auto [num_blocks, block_dims] = get_launch_args(out, large);
   encoder.add_kernel_node(kernel, num_blocks, block_dims, args.args());
 }
 
@@ -229,7 +229,7 @@ void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
   }
   encoder.set_output_array(out);
   auto kernel = mod.get_kernel(kernel_name);
-  auto [num_blocks, block_dims] = get_launch_args(kernel, upd, large);
+  auto [num_blocks, block_dims] = get_launch_args(upd, large);
   encoder.add_kernel_node(kernel, num_blocks, block_dims, args.args());
 }
 
@@ -317,7 +317,7 @@ void GatherAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
   }
   encoder.set_output_array(out);
   auto kernel = mod.get_kernel(kernel_name);
-  auto [num_blocks, block_dims] = get_launch_args(kernel, idx, large);
+  auto [num_blocks, block_dims] = get_launch_args(idx, large);
   encoder.add_kernel_node(kernel, num_blocks, block_dims, args.args());
 }
 
@@ -421,7 +421,7 @@ void ScatterAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
   }
   encoder.set_output_array(out);
   auto kernel = mod.get_kernel(kernel_name);
-  auto [num_blocks, block_dims] = get_launch_args(kernel, idx, large);
+  auto [num_blocks, block_dims] = get_launch_args(idx, large);
   encoder.add_kernel_node(kernel, num_blocks, block_dims, args.args());
 }
 

mlx/backend/cuda/jit_module.cpp

@@ -9,7 +9,6 @@
 #include <cstdlib>
 #include <filesystem>
 #include <fstream>
-#include <unordered_map>
 
 #include <fmt/format.h>
 #include <nvrtc.h>
@@ -330,11 +329,16 @@ CUfunction JitModule::get_kernel(const std::string& kernel_name) {
   return it->second;
 }
 
+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) {
-  static std::unordered_map<std::string, JitModule> map;
+  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).first;

mlx/backend/cuda/jit_module.h

@@ -99,6 +99,8 @@ class JitModule {
   std::unordered_map<std::string, CUfunction> 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.cu

@@ -30,4 +30,25 @@ std::pair<dim3, dim3> get_grid_and_block(int dim0, int dim1, int dim2) {
   return std::make_pair(dim3(gx, gy, gz), dim3(bx, by, bz));
 }
 
+std::tuple<dim3, uint> get_launch_args(
+    size_t size,
+    const Shape& shape,
+    const Strides& strides,
+    bool large,
+    int work_per_thread) {
+  size_t nthreads = cuda::ceil_div(size, work_per_thread);
+  uint block_dim = 1024;
+  if (block_dim > nthreads) {
+    block_dim = nthreads;
+  }
+  dim3 num_blocks;
+  if (large) {
+    num_blocks = get_2d_grid_dims(shape, strides, work_per_thread);
+    num_blocks.x = cuda::ceil_div(num_blocks.x, block_dim);
+  } else {
+    num_blocks.x = cuda::ceil_div(nthreads, block_dim);
+  }
+  return std::make_tuple(num_blocks, block_dim);
+}
+
 } // namespace mlx::core

mlx/backend/cuda/kernel_utils.cuh

@@ -120,53 +120,19 @@ dim3 get_2d_grid_dims(
     size_t divisor);
 std::pair<dim3, dim3> get_grid_and_block(int dim0, int dim1, int dim2);
 
-// Return a block size that achieves maximum potential occupancy for kernel.
-template <typename T>
-inline uint max_occupancy_block_dim(T kernel) {
-  int _, block_dim;
-  if constexpr (std::is_same_v<T, CUfunction>) {
-    CHECK_CUDA_ERROR(
-        cuOccupancyMaxPotentialBlockSize(&_, &block_dim, kernel, 0, 0, 0));
-  } else {
-    CHECK_CUDA_ERROR(
-        cudaOccupancyMaxPotentialBlockSize(&_, &block_dim, kernel));
-  }
-  return block_dim;
-}
-
 // Get the num_blocks and block_dims that maximize occupancy for |kernel|,
 // assuming each thread handles |work_per_thread| elements of |arr|.
-template <typename T>
-inline std::tuple<dim3, uint> get_launch_args(
-    T kernel,
+std::tuple<dim3, uint> get_launch_args(
     size_t size,
     const Shape& shape,
     const Strides& strides,
     bool large,
-    int work_per_thread = 1) {
-  size_t nthreads = cuda::ceil_div(size, work_per_thread);
-  uint block_dim = max_occupancy_block_dim(kernel);
-  if (block_dim > nthreads) {
-    block_dim = nthreads;
-  }
-  dim3 num_blocks;
-  if (large) {
-    num_blocks = get_2d_grid_dims(shape, strides, work_per_thread);
-    num_blocks.x = cuda::ceil_div(num_blocks.x, block_dim);
-  } else {
-    num_blocks.x = cuda::ceil_div(nthreads, block_dim);
-  }
-  return std::make_tuple(num_blocks, block_dim);
-}
+    int work_per_thread = 1);
 
-template <typename T>
-inline std::tuple<dim3, uint> get_launch_args(
-    T kernel,
-    const array& arr,
-    bool large,
-    int work_per_thread = 1) {
+inline std::tuple<dim3, uint>
+get_launch_args(const array& arr, bool large, int work_per_thread = 1) {
   return get_launch_args(
-      kernel, arr.size(), arr.shape(), arr.strides(), large, work_per_thread);
+      arr.size(), arr.shape(), arr.strides(), large, work_per_thread);
 }
 
 } // namespace mlx::core

mlx/backend/cuda/layer_norm.cu

@@ -10,8 +10,6 @@
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
 #include <nvtx3/nvtx3.hpp>
-#include <cub/block/block_load.cuh>
-#include <cub/block/block_reduce.cuh>
 
 namespace mlx::core {
 
@@ -74,9 +72,11 @@ __global__ void layer_norm(
   float sum = 0;
   for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
     auto index = r * BLOCK_DIM + block.thread_rank();
-    T xn[N_READS] = {};
-    cub::LoadDirectBlocked(index, x, xn, axis_size);
-    sum += static_cast<float>(cub::ThreadReduce(xn, cuda::std::plus<>{}));
+    auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      sum += static_cast<float>(xn[i]);
+    }
   }
   sum = BlockReduceT{block, temp}.Sum(sum);
 
@@ -87,11 +87,18 @@ __global__ void layer_norm(
   float normalizer = 0;
   for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
     auto index = r * BLOCK_DIM + block.thread_rank();
-    T xn[N_READS];
-    cub::LoadDirectBlocked(index, x, xn, axis_size, mean);
-    for (int i = 0; i < N_READS; ++i) {
-      float t = static_cast<float>(xn[i]) - mean;
-      normalizer += t * t;
+    if ((index + 1) * N_READS <= axis_size) {
+      auto xn = load_vector<N_READS>(x, index);
+#pragma unroll
+      for (int i = 0; i < N_READS; ++i) {
+        float t = static_cast<float>(xn[i]) - mean;
+        normalizer += t * t;
+      }
+    } else {
+      for (int i = index * N_READS; i < axis_size; ++i) {
+        float t = static_cast<float>(x[i]) - mean;
+        normalizer += t * t;
+      }
     }
   }
   normalizer = BlockReduceT{block, temp}.Sum(normalizer);
@@ -100,17 +107,15 @@ __global__ void layer_norm(
   // Outputs.
   for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
     auto index = r * BLOCK_DIM + block.thread_rank();
-    T xn[N_READS];
-    T wn[N_READS];
-    T bn[N_READS];
-    cub::LoadDirectBlocked(index, x, xn, axis_size);
-    cub::LoadDirectBlocked(index, StridedIterator(w, w_stride), wn, axis_size);
-    cub::LoadDirectBlocked(index, StridedIterator(b, b_stride), bn, axis_size);
+    auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+    auto wn = load_vector<N_READS>(w, index, axis_size, w_stride, T(0));
+    auto bn = load_vector<N_READS>(b, index, axis_size, b_stride, T(0));
+#pragma unroll
     for (int i = 0; i < N_READS; ++i) {
       float norm = (static_cast<float>(xn[i]) - mean) * normalizer;
       xn[i] = wn[i] * static_cast<T>(norm) + bn[i];
     }
-    cub::StoreDirectBlocked(index, out, xn, axis_size);
+    store_vector<N_READS>(out, index, xn, axis_size);
   }
 }
 
@@ -143,9 +148,11 @@ __global__ void layer_norm_vjp(
   float sum = 0;
   for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
     auto index = r * BLOCK_DIM + block.thread_rank();
-    T xn[N_READS] = {};
-    cub::LoadDirectBlocked(index, x, xn, axis_size);
-    sum += static_cast<float>(cub::ThreadReduce(xn, cuda::std::plus<>{}));
+    auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      sum += static_cast<float>(xn[i]);
+    }
   }
   sum = BlockReduceF{block, temp.f}.Sum(sum);
 
@@ -155,19 +162,28 @@ __global__ void layer_norm_vjp(
   // Normalizer.
   float3 factors = {};
   for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
-    T xn[N_READS];
-    T wn[N_READS] = {};
-    T gn[N_READS] = {};
     auto index = r * BLOCK_DIM + block.thread_rank();
-    cub::LoadDirectBlocked(index, x, xn, axis_size, mean);
-    cub::LoadDirectBlocked(index, g, gn, axis_size);
-    cub::LoadDirectBlocked(index, StridedIterator(w, w_stride), wn, axis_size);
-    for (int i = 0; i < N_READS; i++) {
-      float t = static_cast<float>(xn[i]) - mean;
-      float wi = wn[i];
-      float gi = gn[i];
-      float wg = wi * gi;
-      factors = plus_f3(factors, {wg, wg * t, t * t});
+    auto gn = load_vector<N_READS>(g, index, axis_size, T(0));
+    auto wn = load_vector<N_READS>(w, index, axis_size, w_stride, T(0));
+
+    if ((index + 1) * N_READS <= axis_size) {
+      auto xn = load_vector<N_READS>(x, index);
+#pragma unroll
+      for (int i = 0; i < N_READS; ++i) {
+        float t = static_cast<float>(xn[i]) - mean;
+        float wi = wn[i];
+        float gi = gn[i];
+        float wg = wi * gi;
+        factors = plus_f3(factors, {wg, wg * t, t * t});
+      }
+    } else {
+      for (int i = index * N_READS; i < axis_size; ++i) {
+        float t = static_cast<float>(x[i]) - mean;
+        float wi = wn[i];
+        float gi = gn[i];
+        float wg = wi * gi;
+        factors = plus_f3(factors, {wg, wg * t, t * t});
+      }
     }
   }
   factors = BlockReduceF3{block, temp.f3}.Reduce(factors, plus_f3, {});
@@ -179,12 +195,10 @@ __global__ void layer_norm_vjp(
   // Outputs.
   for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
     auto index = r * BLOCK_DIM + block.thread_rank();
-    T xn[N_READS];
-    T wn[N_READS];
-    T gn[N_READS];
-    cub::LoadDirectBlocked(index, x, xn, axis_size);
-    cub::LoadDirectBlocked(index, g, gn, axis_size);
-    cub::LoadDirectBlocked(index, StridedIterator(w, w_stride), wn, axis_size);
+    auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+    auto gn = load_vector<N_READS>(g, index, axis_size, T(0));
+    auto wn = load_vector<N_READS>(w, index, axis_size, w_stride, T(0));
+
     for (int i = 0; i < N_READS; i++) {
       float xi = (static_cast<float>(xn[i]) - mean) * normalizer;
       float wi = wn[i];
@@ -194,9 +208,9 @@ __global__ void layer_norm_vjp(
         wn[i] = gi * xi;
       }
     }
-    cub::StoreDirectBlocked(index, gx, xn, axis_size);
+    store_vector<N_READS>(gx, index, xn, axis_size);
     if constexpr (HAS_W) {
-      cub::StoreDirectBlocked(index, gw, wn, axis_size);
+      store_vector<N_READS>(gw, index, wn, axis_size);
     }
   }
 }
@@ -257,9 +271,9 @@ void LayerNorm::eval_gpu(
   encoder.set_input_array(b);
   encoder.set_output_array(out);
   dispatch_float_types(out.dtype(), "layernorm", [&](auto type_tag) {
-    constexpr uint32_t N_READS = 4;
+    using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    constexpr int N_READS = 16 / sizeof(DataType);
     dispatch_block_dim(cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) {
-      using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
       auto kernel = cu::layer_norm<DataType, block_dim(), N_READS>;
       encoder.add_kernel_node(
           kernel,
@@ -364,10 +378,10 @@ void LayerNormVJP::eval_gpu(
   encoder.set_output_array(gw_temp);
   dispatch_float_types(gx.dtype(), "layernorm_vjp", [&](auto type_tag) {
     dispatch_bool(has_w, [&](auto has_w_constant) {
-      constexpr int N_READS = 4;
+      using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+      constexpr int N_READS = 16 / sizeof(DataType);
       dispatch_block_dim(
           cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) {
-            using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
             auto kernel = cu::layer_norm_vjp<
                 DataType,
                 has_w_constant.value,

mlx/backend/cuda/logsumexp.cu

@@ -43,20 +43,19 @@ __global__ void logsumexp(const T* in, T* out, int axis_size) {
   AccT maxval = Limits<AccT>::finite_min();
   AccT normalizer = 0;
   for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); r++) {
-    AccT vals[N_READS];
-    cub::LoadDirectBlocked(
-        r * BLOCK_DIM + block.thread_rank(),
-        make_cast_iterator<AccT>(in),
-        vals,
-        axis_size,
-        Limits<AccT>::min());
+    auto index = r * BLOCK_DIM + block.thread_rank();
+    auto vals = load_vector<N_READS>(in, index, axis_size, Limits<T>::min());
     prevmax = maxval;
-    maxval = max_op(maxval, cub::ThreadReduce(vals, max_op));
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      maxval = max_op(maxval, static_cast<AccT>(vals[i]));
+    }
     // Online normalizer calculation for softmax:
     // https://github.com/NVIDIA/online-softmax
     normalizer = normalizer * softmax_exp(prevmax - maxval);
     for (int i = 0; i < N_READS; i++) {
-      normalizer = normalizer + softmax_exp(vals[i] - maxval);
+      normalizer =
+          normalizer + softmax_exp(static_cast<AccT>(vals[i]) - maxval);
     }
   }
 
@@ -143,9 +142,9 @@ void LogSumExp::eval_gpu(const std::vector<array>& inputs, array& out) {
   encoder.set_input_array(in);
   encoder.set_output_array(out);
   dispatch_float_types(out.dtype(), "logsumexp", [&](auto type_tag) {
-    constexpr int N_READS = 4;
+    using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    constexpr int N_READS = 16 / sizeof(DataType);
     dispatch_block_dim(cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) {
-      using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
       auto kernel = cu::logsumexp<DataType, float, block_dim(), N_READS>;
       encoder.add_kernel_node(
           kernel,

mlx/backend/cuda/primitives.cpp

@@ -1,47 +1,11 @@
 // Copyright © 2025 Apple Inc.
 
-#include "mlx/backend/cuda/device.h"
-#include "mlx/backend/cuda/device/arange.cuh"
-#include "mlx/backend/cuda/device/fp16_math.cuh"
-#include "mlx/backend/cuda/kernel_utils.cuh"
 #include "mlx/distributed/primitives.h"
-#include "mlx/dtype_utils.h"
 #include "mlx/fast_primitives.h"
 #include "mlx/primitives.h"
 
-#include <nvtx3/nvtx3.hpp>
-#include <thrust/device_ptr.h>
-#include <thrust/transform.h>
-
-#include <cassert>
-
 namespace mlx::core {
 
-void Arange::eval_gpu(const std::vector<array>& inputs, array& out) {
-  nvtx3::scoped_range r("Arange::eval_gpu");
-  assert(inputs.size() == 0);
-  out.set_data(allocator::malloc(out.nbytes()));
-  if (out.size() == 0) {
-    return;
-  }
-  auto& encoder = cu::get_command_encoder(stream());
-  encoder.set_output_array(out);
-  auto capture = encoder.capture_context();
-  dispatch_int_float_types(out.dtype(), "Arange", [&](auto type_tag) {
-    using CTYPE = MLX_GET_TYPE(type_tag);
-    using OutType = cuda_type_t<CTYPE>;
-    CTYPE step =
-        static_cast<CTYPE>(start_ + step_) - static_cast<CTYPE>(start_);
-    thrust::transform(
-        cu::thrust_policy(encoder.stream()),
-        thrust::counting_iterator<uint32_t>(0),
-        thrust::counting_iterator<uint32_t>(out.data_size()),
-        thrust::device_pointer_cast(out.data<OutType>()),
-        cu::Arange<OutType>{
-            static_cast<OutType>(start_), static_cast<OutType>(step)});
-  });
-}
-
 bool fast::ScaledDotProductAttention::use_fallback(
     const array& q,
     const array& k,

mlx/backend/cuda/quantized/affine_quantize.cu

@@ -2,30 +2,17 @@
 
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/kernel_utils.cuh"
-#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/cuda/quantized/quantized_utils.cuh"
 #include "mlx/dtype_utils.h"
-#include "mlx/fast_primitives.h"
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
-#include <nvtx3/nvtx3.hpp>
 
 namespace mlx::core {
 namespace cu {
 
 namespace cg = cooperative_groups;
 
-template <int bits, int wsize = 8>
-inline constexpr __device__ short get_pack_factor() {
-  return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits);
-}
-
-template <int bits, int wsize = 8>
-inline constexpr __device__ short get_bytes_per_pack() {
-  constexpr int power_of_2_bits = (bits & (bits - 1)) == 0;
-  return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3);
-}
-
 template <typename T, int group_size, int bits>
 __global__ void
 affine_quantize(const T* w, uint8_t* out, T* scales, T* biases, size_t size) {
@@ -240,144 +227,100 @@ __global__ void affine_dequantize(
 }
 
 } // namespace cu
-namespace {
 
-inline array ensure_row_contiguous(
-    const array& x,
+void affine_quantize(
+    const array& w,
+    array& wq,
+    array& scales,
+    array& biases,
+    int group_size_,
+    int bits_,
     cu::CommandEncoder& enc,
     const Stream& s) {
-  if (!x.flags().row_contiguous) {
-    array x_copy = contiguous_copy_gpu(x, s);
-    enc.add_temporary(x_copy);
-    return x_copy;
-  } else {
-    return x;
-  }
-}
-
-} // namespace
-
-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 fast::AffineQuantize::eval_gpu(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
-  auto& w_pre = inputs[0];
-  auto& out = outputs[0];
-  out.set_data(allocator::malloc(out.nbytes()));
-
-  auto& s = stream();
-  auto& d = cu::device(s.device);
-  auto& enc = d.get_command_encoder(s);
-
-  auto w = ensure_row_contiguous(w_pre, enc, s);
-  enc.set_input_array(w);
-  if (dequantize_) {
-    auto scales = ensure_row_contiguous(inputs[1], enc, s);
-    auto biases = ensure_row_contiguous(inputs[2], enc, s);
-    enc.set_input_array(scales);
-    enc.set_input_array(biases);
-    enc.set_output_array(out);
-  } else {
-    auto& scales = outputs[1];
-    auto& biases = outputs[2];
-    scales.set_data(allocator::malloc(scales.nbytes()));
-    biases.set_data(allocator::malloc(biases.nbytes()));
-    enc.set_output_array(out);
-    enc.set_output_array(scales);
-    enc.set_output_array(biases);
-  }
-
-  auto dtype = dequantize_ ? outputs[0].dtype() : inputs[0].dtype();
-
-  // Treat uint32 as uint8 in kernel
-  int uint8_per_uint32 = 4;
-  int packs_per_int = (bits_ == 3 || bits_ == 5) ? 8
-      : bits_ == 6                               ? 4
-                                                 : 8 / bits_;
-  int per_thread = dequantize_ ? packs_per_int : group_size_ / WARP_SIZE;
-  size_t size =
-      dequantize_ ? out.size() / packs_per_int : w.size() / per_thread;
+  // Calculate the number of elements per thread
+  int per_thread = group_size_ / WARP_SIZE;
+  size_t size = w.size() / per_thread;
 
+  // Calculate the thread grid that we need to launch
   bool large = size > UINT_MAX;
   auto grid_shape = w.shape();
+  grid_shape.back() /= per_thread;
 
-  if (dequantize_) {
-    grid_shape.back() *= uint8_per_uint32;
-  } else {
-    grid_shape.back() /= per_thread;
-  }
-
-  dispatch_float_types(dtype, "affine_quantize", [&](auto type_tag) {
+  enc.set_input_array(w);
+  enc.set_output_array(wq);
+  enc.set_output_array(scales);
+  enc.set_output_array(biases);
+  dispatch_float_types(w.dtype(), "affine_quantize", [&](auto type_tag) {
     dispatch_groups(group_size_, [&](auto group_size) {
       dispatch_bits(bits_, [&](auto bits) {
-        using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
-        if (dequantize_) {
-          auto kernel =
-              cu::affine_dequantize<DataType, group_size.value, bits.value>;
-          auto [num_blocks, block_dims] =
-              get_launch_args(kernel, size, grid_shape, w.strides(), large);
-          enc.add_kernel_node(
-              kernel,
-              num_blocks,
-              block_dims,
-              w.data<uint8_t>(),
-              inputs[1].data<DataType>(),
-              inputs[2].data<DataType>(),
-              out.data<DataType>(),
-              out.size());
-        } else {
-          auto kernel =
-              cu::affine_quantize<DataType, group_size.value, bits.value>;
-          auto [num_blocks, block_dims] =
-              get_launch_args(kernel, size, grid_shape, w.strides(), large);
-          enc.add_kernel_node(
-              kernel,
-              num_blocks,
-              block_dims,
-              w.data<DataType>(),
-              out.data<uint8_t>(),
-              outputs[1].data<DataType>(),
-              outputs[2].data<DataType>(),
-              w.size());
-        }
+        using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+        auto kernel = cu::affine_quantize<T, group_size.value, bits.value>;
+        auto [num_blocks, block_dims] =
+            get_launch_args(size, grid_shape, w.strides(), large);
+        enc.add_kernel_node(
+            kernel,
+            num_blocks,
+            block_dims,
+            w.data<T>(),
+            wq.data<uint8_t>(),
+            scales.data<T>(),
+            biases.data<T>(),
+            w.size());
+      });
+    });
+  });
+}
+
+void affine_dequantize(
+    const array& wq,
+    const array& scales,
+    const array& biases,
+    array& w,
+    int group_size_,
+    int bits_,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  // Calculate how many numbers we pack together. For 2, 4, 8 bits we pack in
+  // one uint8, for 3, 6 in 3 uint8 and for 5 in 5 uint8.
+  constexpr int uint8_per_uint32 = 4;
+  int packs_per_int;
+  switch (bits_) {
+    case 3:
+    case 5:
+      packs_per_int = 8;
+      break;
+    case 6:
+      packs_per_int = 4;
+      break;
+    default:
+      packs_per_int = 8 / bits_;
+  }
+
+  size_t size = w.size() / packs_per_int;
+  bool large = size > UINT_MAX;
+  auto grid_shape = w.shape();
+  grid_shape.back() *= uint8_per_uint32;
+
+  enc.set_input_array(wq);
+  enc.set_input_array(scales);
+  enc.set_input_array(biases);
+  enc.set_output_array(w);
+  dispatch_float_types(w.dtype(), "affine_quantize", [&](auto type_tag) {
+    dispatch_groups(group_size_, [&](auto group_size) {
+      dispatch_bits(bits_, [&](auto bits) {
+        using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+        auto kernel = cu::affine_dequantize<T, group_size.value, bits.value>;
+        auto [num_blocks, block_dims] =
+            get_launch_args(size, grid_shape, w.strides(), large);
+        enc.add_kernel_node(
+            kernel,
+            num_blocks,
+            block_dims,
+            wq.data<uint8_t>(),
+            scales.data<T>(),
+            biases.data<T>(),
+            w.data<T>(),
+            w.size());
       });
     });
   });

mlx/backend/cuda/quantized/quantized.cpp

@@ -0,0 +1,75 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/quantized/quantized.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/fast_primitives.h"
+
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+namespace {
+
+inline array ensure_row_contiguous(
+    const array& x,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  if (!x.flags().row_contiguous) {
+    array x_copy = contiguous_copy_gpu(x, s);
+    enc.add_temporary(x_copy);
+    return x_copy;
+  } else {
+    return x;
+  }
+}
+
+inline array ensure_row_contiguous_matrix(
+    const array& x,
+    cu::CommandEncoder& enc,
+    const Stream& s) {
+  auto stride_0 = x.strides()[x.ndim() - 2];
+  auto stride_1 = x.strides()[x.ndim() - 1];
+  if (stride_0 == x.shape(-1) && stride_1 == 1) {
+    return x;
+  } else {
+    array x_copy = contiguous_copy_gpu(x, s);
+    enc.add_temporary(x_copy);
+    return x_copy;
+  }
+}
+
+} // namespace
+
+void fast::AffineQuantize::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  nvtx3::scoped_range r("AffineQuantize::eval_gpu");
+  auto& s = stream();
+  auto& d = cu::device(s.device);
+  auto& enc = d.get_command_encoder(s);
+
+  if (dequantize_) {
+    auto wq = ensure_row_contiguous(inputs[0], enc, s);
+    auto scales = ensure_row_contiguous(inputs[1], enc, s);
+    auto biases = ensure_row_contiguous(inputs[2], enc, s);
+    auto& w = outputs[0];
+
+    w.set_data(allocator::malloc(w.nbytes()));
+
+    affine_dequantize(wq, scales, biases, w, group_size_, bits_, enc, s);
+  } else {
+    auto w = ensure_row_contiguous(inputs[0], enc, s);
+    auto& wq = outputs[0];
+    auto& scales = outputs[1];
+    auto& biases = outputs[2];
+
+    wq.set_data(allocator::malloc(wq.nbytes()));
+    scales.set_data(allocator::malloc(scales.nbytes()));
+    biases.set_data(allocator::malloc(biases.nbytes()));
+
+    affine_quantize(w, wq, scales, biases, group_size_, bits_, enc, s);
+  }
+}
+
+} // namespace mlx::core

mlx/backend/cuda/quantized/quantized.h

@@ -0,0 +1,27 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+
+namespace mlx::core {
+
+void affine_quantize(
+    const array& w,
+    array& wq,
+    array& scales,
+    array& biases,
+    int group_size_,
+    int bits_,
+    cu::CommandEncoder& enc,
+    const Stream& s);
+
+void affine_dequantize(
+    const array& wq,
+    const array& scales,
+    const array& biases,
+    array& w,
+    int group_size_,
+    int bits_,
+    cu::CommandEncoder& enc,
+    const Stream& s);
+
+} // namespace mlx::core

mlx/backend/cuda/quantized/quantized_utils.cuh

@@ -0,0 +1,59 @@
+// Copyright © 2025 Apple Inc.
+
+namespace mlx::core {
+
+namespace cu {
+
+template <int bits, int wsize = 8>
+inline constexpr __device__ short get_pack_factor() {
+  return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits);
+}
+
+template <int bits, int wsize = 8>
+inline constexpr __device__ short get_bytes_per_pack() {
+  constexpr int power_of_2_bits = (bits & (bits - 1)) == 0;
+  return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3);
+}
+
+} // 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;
+  }
+}
+
+} // namespace mlx::core

mlx/backend/cuda/reduce.cu

@@ -5,8 +5,6 @@
 #include "mlx/backend/gpu/copy.h"
 
 #include <nvtx3/nvtx3.hpp>
-#include <thrust/device_ptr.h>
-#include <thrust/fill.h>
 
 #include <cassert>
 

mlx/backend/cuda/rms_norm.cu

@@ -10,8 +10,6 @@
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
 #include <nvtx3/nvtx3.hpp>
-#include <cub/block/block_load.cuh>
-#include <cub/block/block_reduce.cuh>
 
 namespace mlx::core {
 
@@ -57,7 +55,7 @@ __global__ void rms_norm(
     const T* w,
     T* out,
     float eps,
-    int32_t axis_size,
+    uint32_t axis_size,
     int64_t w_stride) {
   auto grid = cg::this_grid();
   auto block = cg::this_thread_block();
@@ -72,8 +70,8 @@ __global__ void rms_norm(
   float normalizer = 0;
   for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
     auto index = r * BLOCK_DIM + block.thread_rank();
-    T xn[N_READS];
-    cub::LoadDirectBlocked(index, x, xn, axis_size, cast_to<T>(0));
+    auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+#pragma unroll
     for (int i = 0; i < N_READS; ++i) {
       float t = static_cast<float>(xn[i]);
       normalizer += t * t;
@@ -85,15 +83,14 @@ __global__ void rms_norm(
   // Outputs.
   for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
     auto index = r * BLOCK_DIM + block.thread_rank();
-    T xn[N_READS];
-    T wn[N_READS];
-    cub::LoadDirectBlocked(index, x, xn, axis_size);
-    cub::LoadDirectBlocked(index, StridedIterator(w, w_stride), wn, axis_size);
+    auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+    auto wn = load_vector<N_READS>(w, index, axis_size, w_stride, T(0));
+#pragma unroll
     for (int i = 0; i < N_READS; ++i) {
-      float norm = static_cast<float>(xn[i]) * normalizer;
-      xn[i] = wn[i] * static_cast<T>(norm);
+      float y = static_cast<float>(xn[i]) * normalizer;
+      xn[i] = wn[i] * static_cast<T>(y);
     }
-    cub::StoreDirectBlocked(index, out, xn, axis_size);
+    store_vector<N_READS>(out, index, xn, axis_size);
   }
 }
 
@@ -125,13 +122,10 @@ __global__ void rms_norm_vjp(
   // Normalizer.
   float2 factors = {};
   for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
-    T xn[N_READS];
-    T wn[N_READS] = {};
-    T gn[N_READS] = {};
     auto index = r * BLOCK_DIM + block.thread_rank();
-    cub::LoadDirectBlocked(index, x, xn, axis_size, cast_to<T>(0));
-    cub::LoadDirectBlocked(index, g, gn, axis_size);
-    cub::LoadDirectBlocked(index, StridedIterator(w, w_stride), wn, axis_size);
+    auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+    auto gn = load_vector<N_READS>(g, index, axis_size, T(0));
+    auto wn = load_vector<N_READS>(w, index, axis_size, w_stride, T(0));
     for (int i = 0; i < N_READS; i++) {
       float t = static_cast<float>(xn[i]);
       float wi = wn[i];
@@ -148,12 +142,9 @@ __global__ void rms_norm_vjp(
   // Outputs.
   for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); ++r) {
     auto index = r * BLOCK_DIM + block.thread_rank();
-    T xn[N_READS];
-    T wn[N_READS];
-    T gn[N_READS];
-    cub::LoadDirectBlocked(index, x, xn, axis_size);
-    cub::LoadDirectBlocked(index, g, gn, axis_size);
-    cub::LoadDirectBlocked(index, StridedIterator(w, w_stride), wn, axis_size);
+    auto xn = load_vector<N_READS>(x, index, axis_size, T(0));
+    auto gn = load_vector<N_READS>(g, index, axis_size, T(0));
+    auto wn = load_vector<N_READS>(w, index, axis_size, w_stride, T(0));
     for (int i = 0; i < N_READS; i++) {
       float xi = xn[i];
       float wi = wn[i];
@@ -163,9 +154,9 @@ __global__ void rms_norm_vjp(
         wn[i] = static_cast<T>(gi * xi * normalizer);
       }
     }
-    cub::StoreDirectBlocked(index, gx, xn, axis_size);
+    store_vector<N_READS>(gx, index, xn, axis_size);
     if constexpr (HAS_W) {
-      cub::StoreDirectBlocked(index, gw, wn, axis_size);
+      store_vector<N_READS>(gw, index, wn, axis_size);
     }
   }
 }
@@ -223,9 +214,9 @@ void RMSNorm::eval_gpu(
   encoder.set_input_array(w);
   encoder.set_output_array(out);
   dispatch_float_types(out.dtype(), "rms_norm", [&](auto type_tag) {
-    constexpr uint32_t N_READS = 4;
+    using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    constexpr int N_READS = 16 / sizeof(DataType);
     dispatch_block_dim(cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) {
-      using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
       auto kernel = cu::rms_norm<DataType, block_dim(), N_READS>;
       encoder.add_kernel_node(
           kernel,
@@ -312,11 +303,10 @@ void RMSNormVJP::eval_gpu(
   encoder.set_output_array(gw_temp);
   dispatch_float_types(gx.dtype(), "rms_norm_vjp", [&](auto type_tag) {
     dispatch_bool(has_w, [&](auto has_w_constant) {
-      constexpr int N_READS = 4;
+      using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+      constexpr int N_READS = 16 / sizeof(DataType);
       dispatch_block_dim(
           cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) {
-            using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
-            constexpr int N_READS = 4;
             auto kernel = cu::rms_norm_vjp<
                 DataType,
                 has_w_constant.value,

mlx/backend/cuda/softmax.cu

@@ -11,7 +11,6 @@
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
 #include <nvtx3/nvtx3.hpp>
-#include <cub/block/block_load.cuh>
 
 #include <cassert>
 
@@ -45,20 +44,21 @@ __global__ void softmax(const T* in, T* out, int axis_size) {
   AccT maxval = Limits<AccT>::finite_min();
   AccT normalizer = cast_to<AccT>(0);
   for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); r++) {
-    AccT vals[N_READS];
-    cub::LoadDirectBlocked(
-        r * BLOCK_DIM + block.thread_rank(),
-        make_cast_iterator<AccT>(in),
-        vals,
-        axis_size,
-        Limits<AccT>::min());
+    auto index = r * BLOCK_DIM + block.thread_rank();
+    auto vals = load_vector<N_READS>(in, index, axis_size, Limits<T>::min());
     prevmax = maxval;
-    maxval = max_op(maxval, cub::ThreadReduce(vals, max_op));
+#pragma unroll
+    for (int i = 0; i < N_READS; ++i) {
+      maxval = max_op(maxval, static_cast<AccT>(vals[i]));
+    }
+
     // Online normalizer calculation for softmax:
     // https://github.com/NVIDIA/online-softmax
     normalizer = normalizer * softmax_exp(prevmax - maxval);
+#pragma unroll
     for (int i = 0; i < N_READS; i++) {
-      normalizer = normalizer + softmax_exp(vals[i] - maxval);
+      normalizer =
+          normalizer + softmax_exp(static_cast<AccT>(vals[i]) - maxval);
     }
   }
 
@@ -95,12 +95,11 @@ __global__ void softmax(const T* in, T* out, int axis_size) {
   // Write output.
   for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); r++) {
     auto index = r * BLOCK_DIM + block.thread_rank();
-    T vals[N_READS];
-    cub::LoadDirectBlocked(index, in, vals, axis_size);
+    auto vals = load_vector<N_READS>(in, index, axis_size, T(0));
     for (int i = 0; i < N_READS; i++) {
       vals[i] = softmax_exp(static_cast<AccT>(vals[i]) - maxval) * normalizer;
     }
-    cub::StoreDirectBlocked(index, out, vals, axis_size);
+    store_vector<N_READS>(out, index, vals, axis_size);
   }
 }
 
@@ -141,9 +140,9 @@ void Softmax::eval_gpu(const std::vector<array>& inputs, array& out) {
   encoder.set_input_array(in);
   encoder.set_output_array(out);
   dispatch_float_types(out.dtype(), "softmax", [&](auto type_tag) {
-    constexpr int N_READS = 4;
+    using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    constexpr int N_READS = 16 / sizeof(DataType);
     dispatch_block_dim(cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) {
-      using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
       auto kernel = cu::softmax<DataType, DataType, block_dim(), N_READS>;
       if (precise) {
         kernel = cu::softmax<DataType, float, block_dim(), N_READS>;

mlx/backend/cuda/sort.cu

@@ -13,7 +13,6 @@
 #include <cub/device/device_segmented_sort.cuh>
 
 #include <cassert>
-#include <numeric>
 
 namespace mlx::core {
 
@@ -27,29 +26,6 @@ struct ModOp {
   }
 };
 
-// We can not use any op in eval, make an utility.
-array swapaxes_in_eval(const array& in, int axis1, int axis2) {
-  std::vector<int> axes(in.ndim());
-  std::iota(axes.begin(), axes.end(), 0);
-  std::swap(axes[axis1], axes[axis2]);
-  // TODO: Share the code with Transpose::eval.
-  Shape shape(axes.size());
-  Strides strides(in.ndim());
-  for (size_t ax = 0; ax < axes.size(); ++ax) {
-    shape[ax] = in.shape()[axes[ax]];
-    strides[ax] = in.strides()[axes[ax]];
-  }
-  auto flags = in.flags();
-  if (flags.contiguous) {
-    auto [_, row_contiguous, col_contiguous] = check_contiguity(shape, strides);
-    flags.row_contiguous = row_contiguous;
-    flags.col_contiguous = col_contiguous;
-  }
-  array out(shape, in.dtype(), nullptr, {});
-  out.copy_shared_buffer(in, strides, flags, in.data_size());
-  return out;
-}
-
 struct OffsetTransform {
   int nsort;
 

mlx/backend/cuda/ternary.cu

@@ -32,7 +32,7 @@ ternary_v(const bool* a, const T* b, const T* c, T* out, IdxT size) {
     AlignedVector<T, N_READS> out_vec;
 #pragma unroll
     for (int i = 0; i < N_READS; ++i) {
-      out_vec.val[i] = Op{}(a_vec.val[i], b_vec.val[i], c_vec.val[i]);
+      out_vec[i] = Op{}(a_vec[i], b_vec[i], c_vec[i]);
     }
 
     store_vector<N_READS>(out, index, out_vec);
@@ -125,12 +125,9 @@ void ternary_op_gpu_inplace(
             int ndim = shape.size();
             if (ndim <= 3) {
               dispatch_1_2_3(ndim, [&](auto dims_constant) {
-                auto kernel =
-                    cu::ternary_g_nd<Op, DType, IdxT, dims_constant()>;
-                auto [num_blocks, block_dims] =
-                    get_launch_args(kernel, out, large());
+                auto [num_blocks, block_dims] = get_launch_args(out, large());
                 encoder.add_kernel_node(
-                    kernel,
+                    cu::ternary_g_nd<Op, DType, IdxT, dims_constant()>,
                     num_blocks,
                     block_dims,
                     a.data<bool>(),
@@ -144,11 +141,9 @@ void ternary_op_gpu_inplace(
                     const_param<dims_constant()>(c_strides));
               });
             } else {
-              auto kernel = cu::ternary_g<Op, DType, IdxT>;
-              auto [num_blocks, block_dims] =
-                  get_launch_args(kernel, out, large());
+              auto [num_blocks, block_dims] = get_launch_args(out, large());
               encoder.add_kernel_node(
-                  kernel,
+                  cu::ternary_g<Op, DType, IdxT>,
                   num_blocks,
                   block_dims,
                   a.data<bool>(),
@@ -166,18 +161,11 @@ void ternary_op_gpu_inplace(
     } else {
       dispatch_bool(out.data_size() > UINT32_MAX, [&](auto large) {
         using IdxT = std::conditional_t<large(), int64_t, uint32_t>;
-        // TODO: Choose optimized value based on type size.
-        constexpr int N_READS = 4;
-        auto kernel = cu::ternary_v<Op, DType, IdxT, N_READS>;
+        constexpr int N_READS = 16 / sizeof(DType);
         auto [num_blocks, block_dims] = get_launch_args(
-            kernel,
-            out.data_size(),
-            out.shape(),
-            out.strides(),
-            large(),
-            N_READS);
+            out.data_size(), out.shape(), out.strides(), large(), N_READS);
         encoder.add_kernel_node(
-            kernel,
+            cu::ternary_v<Op, DType, IdxT, N_READS>,
             num_blocks,
             block_dims,
             a.data<bool>(),
@@ -204,7 +192,7 @@ void ternary_op_gpu(
 }
 
 void Select::eval_gpu(const std::vector<array>& inputs, array& out) {
-  nvtx3::scoped_range r("select::eval_gpu");
+  nvtx3::scoped_range r("Select::eval_gpu");
   auto& s = out.primitive().stream();
   ternary_op_gpu<cu::Select>(inputs, out, s);
 }

mlx/backend/cuda/unary.cu

@@ -30,7 +30,7 @@ __global__ void unary_v(const In* in, Out* out, IdxT size) {
     AlignedVector<Out, N_READS> out_vec;
 #pragma unroll
     for (int i = 0; i < N_READS; ++i) {
-      out_vec.val[i] = Op{}(in_vec.val[i]);
+      out_vec[i] = Op{}(in_vec[i]);
     }
 
     store_vector<N_READS>(out, index, out_vec);
@@ -129,16 +129,10 @@ void unary_op_gpu_inplace(
             using IdxT = std::conditional_t<large(), int64_t, uint32_t>;
             // TODO: Choose optimized value based on type size.
             constexpr int N_READS = 4;
-            auto kernel = cu::unary_v<Op, InType, OutType, IdxT, N_READS>;
             auto [num_blocks, block_dims] = get_launch_args(
-                kernel,
-                out.data_size(),
-                out.shape(),
-                out.strides(),
-                large,
-                N_READS);
+                out.data_size(), out.shape(), out.strides(), large, N_READS);
             encoder.add_kernel_node(
-                kernel,
+                cu::unary_v<Op, InType, OutType, IdxT, N_READS>,
                 num_blocks,
                 block_dims,
                 in.data<InType>(),
@@ -147,10 +141,9 @@ void unary_op_gpu_inplace(
           } else {
             using IdxT = std::conditional_t<large(), int64_t, int32_t>;
             auto [shape, strides] = collapse_contiguous_dims(in);
-            auto kernel = cu::unary_g<Op, InType, OutType, IdxT>;
-            auto [num_blocks, block_dims] = get_launch_args(kernel, out, large);
+            auto [num_blocks, block_dims] = get_launch_args(out, large);
             encoder.add_kernel_node(
-                kernel,
+                cu::unary_g<Op, InType, OutType, IdxT>,
                 num_blocks,
                 block_dims,
                 in.data<InType>(),

mlx/backend/gpu/primitives.cpp

@@ -133,6 +133,7 @@ void NumberOfElements::eval_gpu(const std::vector<array>& inputs, array& out) {
 }
 
 void Pad::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("Pad::eval_gpu");
   // Inputs must be base input array and scalar val array
   assert(inputs.size() == 2);
   auto& in = inputs[0];

python/mlx/nn/layers/base.py

@@ -399,41 +399,7 @@ class Module(dict):
         Returns:
             The module instance after updating the submodules.
         """
-
-        def apply(dst, modules):
-            if isinstance(modules, dict):
-                for k in modules:
-                    if k in dst:
-                        current_value = dst[k]
-                        new_value = modules[k]
-                        if self.is_module(current_value) and self.is_module(new_value):
-                            dst[k] = new_value
-                        elif isinstance(current_value, (dict, list)):
-                            apply(current_value, new_value)
-                        elif strict and new_value != {}:
-                            raise ValueError(
-                                f"Received invalid type: {type(new_value).__name__}."
-                            )
-                    elif strict:
-                        raise ValueError(
-                            f'Module does not have sub-module named "{k}".'
-                        )
-            elif isinstance(modules, list):
-                for i in range(len(modules)):
-                    current_value = dst[i]
-                    new_value = modules[i]
-                    if self.is_module(current_value) and self.is_module(new_value):
-                        dst[i] = new_value
-                    elif isinstance(current_value, (dict, list)):
-                        apply(current_value, new_value)
-                    elif strict and new_value != {}:
-                        raise ValueError(
-                            f"Received invalid type: {type(new_value).__name__}."
-                        )
-            elif strict:
-                raise ValueError(f"Received invalid type: {type(modules).__name__}.")
-
-        apply(self, modules)
+        _update_modules(self, modules, strict)
         return self
 
     def apply_to_modules(self, apply_fn: Callable[[str, Module], Any]) -> Module:
@@ -639,6 +605,36 @@ class Module(dict):
         self.apply(lambda x: x.astype(dtype) if predicate(x.dtype) else x)
 
 
+def _update_modules(dst, modules, strict):
+    if isinstance(modules, dict):
+        for k in modules:
+            if k in dst:
+                current_value = dst[k]
+                new_value = modules[k]
+                if Module.is_module(current_value) and Module.is_module(new_value):
+                    dst[k] = new_value
+                elif isinstance(current_value, (dict, list)):
+                    _update_modules(current_value, new_value, strict)
+                elif strict and new_value != {}:
+                    raise ValueError(
+                        f"Received invalid type: {type(new_value).__name__}."
+                    )
+            elif strict:
+                raise ValueError(f'Module does not have sub-module named "{k}".')
+    elif isinstance(modules, list):
+        for i in range(len(modules)):
+            current_value = dst[i]
+            new_value = modules[i]
+            if Module.is_module(current_value) and Module.is_module(new_value):
+                dst[i] = new_value
+            elif isinstance(current_value, (dict, list)):
+                _update_modules(current_value, new_value, strict)
+            elif strict and new_value != {}:
+                raise ValueError(f"Received invalid type: {type(new_value).__name__}.")
+    elif strict:
+        raise ValueError(f"Received invalid type: {type(modules).__name__}.")
+
+
 def _unwrap(model, value_key, value, filter_fn, map_fn, is_leaf_fn):
     if is_leaf_fn(model, value_key, value):
         return map_fn(value)

python/tests/cuda_skip.py

@@ -18,21 +18,12 @@ cuda_skip = {
     "TestConv.test_conv_1d_groups_flipped",
     "TestConv.test_conv_general_flip_grad",
     "TestConv.test_conv_groups_grad",
-    "TestConv.test_torch_conv_1D_grad",
     "TestConv.test_torch_conv_2D",
-    "TestConv.test_torch_conv_2D_grad",
-    "TestConv.test_torch_conv_3D_grad",
     "TestConv.test_torch_conv_depthwise",
     "TestConv.test_torch_conv_general",
-    "TestConvTranspose.test_torch_conv_tranpose_1d_output_padding",
-    "TestConvTranspose.test_torch_conv_transpose_1D",
     "TestConvTranspose.test_torch_conv_transpose_1D_grad",
-    "TestConvTranspose.test_torch_conv_transpose_2D",
     "TestConvTranspose.test_torch_conv_transpose_2D_grad",
-    "TestConvTranspose.test_torch_conv_transpose_2d_output_padding",
-    "TestConvTranspose.test_torch_conv_transpose_3D",
     "TestConvTranspose.test_torch_conv_transpose_3D_grad",
-    "TestConvTranspose.test_torch_conv_transpose_3d_output_padding",
     # FFTs NYI
     "TestFFT.test_fft",
     "TestFFT.test_fft_big_powers_of_two",
@@ -74,4 +65,5 @@ cuda_skip = {
     "TestQuantized.test_small_matrix",
     "TestQuantized.test_throw",
     "TestQuantized.test_vjp_scales_biases",
+    "TestExportImport.test_export_quantized_model",
 }

python/tests/test_blas.py

@@ -47,7 +47,7 @@ class TestBlas(mlx_tests.MLXTestCase):
             self.assertTrue(np.allclose(out_mlx, out_npy.astype(np_dtype), atol=1e-5))
 
     def test_matmul_unaligned(self):
-        if not mx.metal.is_available():
+        if not mx.is_available(mx.gpu):
             return
 
         for dtype in self.dtypes:
@@ -61,8 +61,15 @@ class TestBlas(mlx_tests.MLXTestCase):
                     shape_b = (dim + p, dim + p)
                     self.__gemm_test(shape_a, shape_b, np_dtype)
 
+    def test_matvec_unaligned(self):
+        a = mx.random.normal(shape=(4, 128))
+        b = mx.random.normal(shape=(129,))[1:]
+        out = a @ b
+        np_out = np.array(a) @ np.array(b)
+        self.assertTrue(np.allclose(out, np_out))
+
     def test_matmul_shapes(self):
-        if not mx.metal.is_available():
+        if not mx.is_available(mx.gpu):
             return
 
         shapes = [
@@ -1274,7 +1281,7 @@ class TestBlas(mlx_tests.MLXTestCase):
     def test_gemv_gemm_same_precision(self):
         mx.random.seed(0)
         N = 256
-        if mx.metal.is_available():
+        if mx.is_available(mx.gpu):
             t = mx.bfloat16
             a = mx.random.normal([1, N]).astype(t)
             b = mx.concatenate([a, a], axis=0).astype(t)

python/tests/test_export_import.py

@@ -346,6 +346,46 @@ class TestExportImport(mlx_tests.MLXTestCase):
         expected = forward(input_data)
         self.assertTrue(mx.allclose(expected, out))
 
+    def test_export_control_flow(self):
+
+        def fun(x, y):
+            if y.shape[0] <= 2:
+                return x + y
+            else:
+                return x + 2 * y
+
+        for y in (mx.array([1, 2, 3]), mx.array([1, 2])):
+            for shapeless in (True, False):
+                with self.subTest(y=y, shapeless=shapeless):
+                    x = mx.array(1)
+                    export_path = os.path.join(self.test_dir, "control_flow.mlxfn")
+                    mx.export_function(export_path, fun, x, y, shapeless=shapeless)
+
+                    imported_fn = mx.import_function(export_path)
+                    self.assertTrue(mx.array_equal(imported_fn(x, y)[0], fun(x, y)))
+
+    def test_export_quantized_model(self):
+        for shapeless in (True, False):
+            with self.subTest(shapeless=shapeless):
+                model = nn.Sequential(
+                    nn.Linear(1024, 512), nn.ReLU(), nn.Linear(512, 1024)
+                )
+                model.eval()
+                mx.eval(model.parameters())
+                input_data = mx.ones(shape=(512, 1024))
+                nn.quantize(model)
+                self.assertTrue(isinstance(model.layers[0], nn.QuantizedLinear))
+                self.assertTrue(isinstance(model.layers[2], nn.QuantizedLinear))
+                mx.eval(model.parameters())
+
+                export_path = os.path.join(self.test_dir, "quantized_linear.mlxfn")
+                mx.export_function(export_path, model, input_data, shapeless=shapeless)
+
+                imported_fn = mx.import_function(export_path)
+                self.assertTrue(
+                    mx.array_equal(imported_fn(input_data)[0], model(input_data))
+                )
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()

python/tests/test_nn.py

@@ -279,6 +279,23 @@ class TestBase(mlx_tests.MLXTestCase):
         del m.weight
         self.assertFalse(hasattr(m, "weight"))
 
+    def test_circular_leaks(self):
+        y = mx.random.uniform(1)
+        mx.eval(y)
+
+        def make_and_update():
+            model = nn.Linear(1024, 512)
+            mx.eval(model.parameters())
+            leaves = {}
+            model.update_modules(leaves)
+
+        mx.synchronize()
+        pre = mx.get_active_memory()
+        make_and_update()
+        mx.synchronize()
+        post = mx.get_active_memory()
+        self.assertEqual(pre, post)
+
 
 class TestLayers(mlx_tests.MLXTestCase):
     def test_identity(self):

python/tests/test_ops.py

@@ -3049,6 +3049,25 @@ class TestOps(mlx_tests.MLXTestCase):
         out = mx.power(mx.array(0j), float("nan"))
         self.assertTrue(mx.isnan(out))
 
+    def test_irregular_alignments(self):
+        # Unaligned unary op
+        a = mx.ones((64, 1))
+        b = -a[1:]
+        self.assertTrue(mx.all(b == -1.0))
+
+        # Unaligned binary op
+        a = mx.ones((64, 1))
+        b = a[1:]
+        c = b + b
+        self.assertTrue(mx.all(c == 2.0))
+
+        # Unaligned ternary op
+        a = mx.ones((64, 1))
+        b = mx.zeros((63, 1))
+        c = mx.ones((63, 1)).astype(mx.bool_)
+        d = mx.where(c, a[1:], b)
+        self.assertTrue(mx.all(d == 1.0))
+
 
 class TestBroadcast(mlx_tests.MLXTestCase):
     def test_broadcast_shapes(self):

setup.py

@@ -44,6 +44,8 @@ def get_version():
 
 
 build_stage = int(os.environ.get("MLX_BUILD_STAGE", 0))
+build_macos = platform.system() == "Darwin"
+build_cuda = "MLX_BUILD_CUDA=ON" in os.environ.get("CMAKE_ARGS", "")
 
 
 # A CMakeExtension needs a sourcedir instead of a file list.
@@ -85,6 +87,11 @@ class CMakeBuild(build_ext):
             "-DMLX_BUILD_EXAMPLES=OFF",
             f"-DMLX_PYTHON_BINDINGS_OUTPUT_DIRECTORY={extdir}{os.sep}",
         ]
+        if build_stage == 2 and build_cuda:
+            # Last arch is always real and virtual for forward-compatibility
+            cuda_archs = ";".join(("70-real", "80-real", "90-real", "100-real", "120"))
+            cmake_args += [f"-DMLX_CUDA_ARCHITECTURES={cuda_archs}"]
+
         # Some generators require explcitly passing config when building.
         build_args = ["--config", cfg]
         # Adding CMake arguments set as environment variable
@@ -95,7 +102,7 @@ class CMakeBuild(build_ext):
         # Pass version to C++
         cmake_args += [f"-DMLX_VERSION={self.distribution.get_version()}"]  # type: ignore[attr-defined]
 
-        if platform.system() == "Darwin":
+        if build_macos:
             # Cross-compile support for macOS - respect ARCHFLAGS if set
             archs = re.findall(r"-arch (\S+)", os.environ.get("ARCHFLAGS", ""))
             if archs:
@@ -113,6 +120,9 @@ class CMakeBuild(build_ext):
         if "CMAKE_BUILD_PARALLEL_LEVEL" not in os.environ:
             build_args += [f"-j{os.cpu_count()}"]
 
+        # Avoid cache miss when building from temporary dirs.
+        os.environ["CCACHE_BASEDIR"] = os.path.abspath(self.build_temp)
+
         subprocess.run(
             ["cmake", ext.sourcedir, *cmake_args], cwd=build_temp, check=True
         )
@@ -202,9 +212,6 @@ if __name__ == "__main__":
         ],
     )
 
-    build_macos = platform.system() == "Darwin"
-    build_cuda = "MLX_BUILD_CUDA=ON" in os.environ.get("CMAKE_ARGS", "")
-
     version = get_version()
 
     _setup = partial(