split logsumexp

CMakeLists.txt

@@ -34,7 +34,6 @@ option(MLX_BUILD_BENCHMARKS "Build benchmarks for mlx" OFF)
 option(MLX_BUILD_PYTHON_BINDINGS "Build python bindings for mlx" OFF)
 option(MLX_BUILD_METAL "Build metal backend" ON)
 option(MLX_BUILD_CPU "Build cpu backend" ON)
-option(MLX_BUILD_CUDA "Build cuda backend" OFF)
 option(MLX_METAL_DEBUG "Enhance metal debug workflow" OFF)
 option(MLX_ENABLE_X64_MAC "Enable building for x64 macOS" OFF)
 option(MLX_BUILD_GGUF "Include support for GGUF format" ON)
@@ -84,10 +83,6 @@ if(MLX_BUILD_METAL)
   set(QUARTZ_LIB "-framework QuartzCore")
 endif()
 
-if(MLX_BUILD_CUDA)
-  enable_language(CUDA)
-endif()
-
 if(MLX_BUILD_METAL AND NOT METAL_LIB)
   message(STATUS "Metal not found. Unable to build GPU")
   set(MLX_BUILD_METAL OFF)

mlx/CMakeLists.txt

@@ -47,18 +47,10 @@ add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/distributed)
 add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/io)
 
 if(MLX_BUILD_METAL)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/gpu)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/metal)
 else()
   target_sources(mlx
                  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/backend/metal/no_metal.cpp)
-endif()
-
-if(MLX_BUILD_CUDA)
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/cuda)
-endif()
-
-if(MLX_BUILD_METAL OR MLX_BUILD_CUDA)
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/gpu)
-else()
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/no_gpu)
 endif()

mlx/backend/common/CMakeLists.txt

@@ -6,5 +6,4 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/load.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/slicing.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/transpose.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp)

mlx/backend/common/common.cpp

@@ -2,7 +2,6 @@
 #include <cassert>
 
 #include "mlx/backend/common/broadcasting.h"
-#include "mlx/backend/common/transpose.h"
 #include "mlx/backend/common/utils.h"
 #include "mlx/primitives.h"
 
@@ -20,19 +19,26 @@ void AsStrided::eval(const std::vector<array>& inputs, array& out) {
         "AsStrided must be used with row contiguous arrays only.");
   }
 
-  // Calculate the contiguity based on the given shape and strides
-  auto [ds, rc, cc] = check_contiguity(shape_, strides_);
+  // Compute the flags given the shape and strides
+  bool row_contiguous = true, col_contiguous = true;
+  size_t r = 1, c = 1;
+  for (int i = strides_.size() - 1, j = 0; i >= 0; i--, j++) {
+    row_contiguous &= (r == strides_[i]) || (shape_[i] == 1);
+    col_contiguous &= (c == strides_[j]) || (shape_[j] == 1);
+    r *= shape_[i];
+    c *= shape_[j];
+  }
   auto flags = in.flags();
-
   // TODO: Compute the contiguous flag in a better way cause now we are
   //       unnecessarily strict.
-  flags.contiguous = rc || cc;
-  flags.row_contiguous = rc;
-  flags.col_contiguous = cc;
+  flags.contiguous = row_contiguous || col_contiguous;
+  flags.row_contiguous = row_contiguous;
+  flags.col_contiguous = col_contiguous;
 
-  // There is no easy way to compute the actual data size so we use out.size()
-  // when the array is not contiguous.
-  size_t data_size = flags.contiguous ? ds : out.size();
+  // There is no easy way to compute the actual data size so we use out.size().
+  // The contiguous flag will almost certainly not be set so no code should
+  // rely on data_size anyway.
+  size_t data_size = out.size();
 
   return out.copy_shared_buffer(in, strides_, flags, data_size, offset_);
 }
@@ -264,7 +270,36 @@ void StopGradient::eval(const std::vector<array>& inputs, array& out) {
 
 void Transpose::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
-  transpose(inputs[0], out, axes_);
+  Strides out_strides(out.ndim());
+  auto& in = inputs[0];
+  for (int ax = 0; ax < axes_.size(); ++ax) {
+    out_strides[ax] = in.strides()[axes_[ax]];
+  }
+
+  // Conditions for {row/col}_contiguous
+  // - array must be contiguous (no gaps)
+  // - underlying buffer size should have the same size as the array
+  // - cumulative product of shapes is equal to the strides (we can ignore axes
+  //   with size == 1)
+  //   - in the forward direction (column contiguous)
+  //   - in the reverse direction (row contiguous)
+  // - vectors are both row and col contiguous (hence if both row/col are
+  //   true, they stay true)
+  auto flags = in.flags();
+  if (flags.contiguous && in.data_size() == in.size()) {
+    int64_t f_stride = 1;
+    int64_t b_stride = 1;
+    flags.col_contiguous = true;
+    flags.row_contiguous = true;
+    for (int i = 0, ri = out.ndim() - 1; i < out.ndim(); ++i, --ri) {
+      flags.col_contiguous &= (out_strides[i] == f_stride || out.shape(i) == 1);
+      f_stride *= out.shape(i);
+      flags.row_contiguous &=
+          (out_strides[ri] == b_stride || out.shape(ri) == 1);
+      b_stride *= out.shape(ri);
+    }
+  }
+  out.copy_shared_buffer(in, out_strides, flags, in.data_size());
 }
 
 } // namespace mlx::core

mlx/backend/common/transpose.cpp

@@ -1,57 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-#include <cassert>
-
-#include "mlx/backend/common/utils.h"
-
-namespace mlx::core {
-
-void transpose(const array& in, array& out, const std::vector<int>& axes) {
-  Strides out_strides(out.ndim());
-  for (int ax = 0; ax < axes.size(); ++ax) {
-    out_strides[ax] = in.strides()[axes[ax]];
-  }
-
-  // Conditions for {row/col}_contiguous
-  // - array must be contiguous (no gaps)
-  // - underlying buffer size should have the same size as the array
-  // - cumulative product of shapes is equal to the strides (we can ignore axes
-  //   with size == 1)
-  //   - in the forward direction (column contiguous)
-  //   - in the reverse direction (row contiguous)
-  // - vectors are both row and col contiguous (hence if both row/col are
-  //   true, they stay true)
-  auto flags = in.flags();
-  if (flags.contiguous && in.data_size() == in.size()) {
-    auto [_, rc, cc] = check_contiguity(out.shape(), out_strides);
-    flags.row_contiguous = rc;
-    flags.col_contiguous = cc;
-  }
-  out.copy_shared_buffer(in, out_strides, flags, in.data_size());
-}
-
-void as_transposed(array& out, const std::vector<int>& axes) {
-  assert(out.data_size() == out.size() && out.flags().contiguous);
-
-  // Calculate the contiguous strides.
-  Strides strides(out.ndim(), 1);
-  for (int i = out.ndim() - 2; i >= 0; i--) {
-    strides[i] = strides[i + 1] * out.shape(i);
-  }
-
-  // Calculate the new strides for transposing.
-  Strides new_strides;
-  new_strides.reserve(out.ndim());
-  for (auto ax : axes) {
-    new_strides.push_back(strides[ax]);
-  }
-
-  auto [ds, rc, cc] = check_contiguity(out.shape(), new_strides);
-  auto flags = out.flags();
-  flags.row_contiguous = rc;
-  flags.col_contiguous = cc;
-
-  out.copy_shared_buffer(out, new_strides, flags, ds);
-}
-
-} // namespace mlx::core

mlx/backend/common/transpose.h

@@ -1,12 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-#pragma once
-
-#include "mlx/array.h"
-
-namespace mlx::core {
-
-void transpose(const array& in, array& out, const std::vector<int>& axes);
-void as_transposed(array& out, const std::vector<int>& axes);
-
-} // namespace mlx::core

mlx/backend/common/utils.h

@@ -132,11 +132,6 @@ struct ContiguousIterator {
 };
 
 inline auto check_contiguity(const Shape& shape, const Strides& strides) {
-  // Conditions for {row/col}_contiguous
-  // - cumulative product of shapes is equal to the strides (we can ignore axes
-  //   with size == 1)
-  //   - in the forward direction (column contiguous)
-  //   - in the reverse direction (row contiguous)
   size_t no_broadcast_data_size = 1;
   int64_t f_stride = 1;
   int64_t b_stride = 1;

mlx/backend/cuda/CMakeLists.txt

@@ -1,57 +0,0 @@
-# Filename rules in cuda backend:
-#
-# * Use .cu/.cuh if code contains device code, and .cpp/.h if not.
-# * Device-only kernel code should be put in kernels/ subdir.
-# * Files in kernels/ subdir should not include files outside.
-target_sources(
-  mlx
-  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/allocator.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/copy.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/eval.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/event.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/fence.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cu
-          ${CMAKE_CURRENT_SOURCE_DIR}/slicing.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/worker.cpp)
-
-target_compile_definitions(mlx PUBLIC MLX_USE_CUDA)
-
-# Enable defining device lambda functions.
-target_compile_options(mlx
-                       PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--extended-lambda>")
-
-# 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
-    "75;80"
-    CACHE STRING "CUDA architectures")
-message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}")
-set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES
-                                     "${MLX_CUDA_ARCHITECTURES}")
-
-# Use fixed version of CCCL.
-FetchContent_Declare(
-  cccl
-  URL "https://github.com/NVIDIA/cccl/releases/download/v2.8.1/cccl-v2.8.1.zip")
-FetchContent_MakeAvailable(cccl)
-target_include_directories(mlx PRIVATE BEFORE "${cccl_SOURCE_DIR}/include")
-
-# Use fixed version of NVTX.
-FetchContent_Declare(
-  nvtx3
-  GIT_REPOSITORY https://github.com/NVIDIA/NVTX.git
-  GIT_TAG v3.1.1
-  GIT_SHALLOW TRUE
-  SOURCE_SUBDIR c EXCLUDE_FROM_ALL)
-FetchContent_MakeAvailable(nvtx3)
-target_link_libraries(mlx PUBLIC $<BUILD_INTERFACE:nvtx3-cpp>)
-
-# Make cuda runtime APIs available in non-cuda files.
-find_package(CUDAToolkit REQUIRED)
-target_include_directories(mlx PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
-
-# Suppress nvcc warnings on MLX headers.
-target_compile_options(mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe
-                                   --diag_suppress=997>)

mlx/backend/cuda/allocator.cpp

@@ -1,154 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/cuda/allocator.h"
-#include "mlx/backend/cuda/utils.h"
-#include "mlx/backend/cuda/worker.h"
-
-#include <cuda_runtime.h>
-#include <fmt/format.h>
-
-#include <cassert>
-
-namespace mlx::core {
-
-namespace cu {
-
-CudaAllocator::CudaAllocator() {
-  // TODO: Set memory limit for multi-device.
-  size_t free, total;
-  CHECK_CUDA_ERROR(cudaMemGetInfo(&free, &total));
-  memory_limit_ = total * 0.8;
-}
-
-Buffer CudaAllocator::malloc(size_t size) {
-  // TODO: Check memory limit.
-  auto* buf = new CudaBuffer{nullptr, size};
-  cudaError_t err = cudaMallocManaged(&buf->data, size);
-  if (err != cudaSuccess && err != cudaErrorMemoryAllocation) {
-    throw std::runtime_error(
-        fmt::format("cudaMallocManaged failed: {}.", cudaGetErrorString(err)));
-  }
-  std::lock_guard lock(mutex_);
-  active_memory_ += size;
-  peak_memory_ = std::max(active_memory_, peak_memory_);
-  return Buffer{buf};
-}
-
-void CudaAllocator::free(Buffer buffer) {
-  auto* buf = static_cast<CudaBuffer*>(buffer.ptr());
-  if (!buf) {
-    return;
-  }
-
-  // If free() is called from a unregistered thread, reschedule the call to
-  // worker.
-  {
-    std::lock_guard lock(worker_mutex_);
-    if (allowed_threads_.count(std::this_thread::get_id()) == 0) {
-      if (!worker_) {
-        worker_.reset(new Worker);
-      }
-      worker_->add_task([buffer]() { allocator().free(buffer); });
-      worker_->end_batch();
-      worker_->commit();
-      return;
-    }
-  }
-
-  size_t size = buf->size;
-  cudaFree(buf->data);
-  delete buf;
-  std::lock_guard lock(mutex_);
-  active_memory_ -= size;
-}
-
-size_t CudaAllocator::size(Buffer buffer) const {
-  auto* buf = static_cast<CudaBuffer*>(buffer.ptr());
-  if (!buf) {
-    return 0;
-  }
-  return buf->size;
-}
-
-void CudaAllocator::register_this_thread() {
-  std::lock_guard lock(worker_mutex_);
-  allowed_threads_.insert(std::this_thread::get_id());
-}
-
-size_t CudaAllocator::get_active_memory() const {
-  return active_memory_;
-}
-
-size_t CudaAllocator::get_peak_memory() const {
-  return peak_memory_;
-}
-
-void CudaAllocator::reset_peak_memory() {
-  std::lock_guard lock(mutex_);
-  peak_memory_ = 0;
-}
-
-size_t CudaAllocator::get_memory_limit() {
-  return memory_limit_;
-}
-
-size_t CudaAllocator::set_memory_limit(size_t limit) {
-  std::lock_guard lock(mutex_);
-  std::swap(limit, memory_limit_);
-  return limit;
-}
-
-CudaAllocator& allocator() {
-  // By creating the |allocator_| on heap, the destructor of CudaAllocator
-  // will not be called on exit and buffers in the cache will be leaked. This
-  // can save some time at program exit.
-  static CudaAllocator* allocator_ = new CudaAllocator;
-  return *allocator_;
-}
-
-} // namespace cu
-
-namespace allocator {
-
-Allocator& allocator() {
-  return cu::allocator();
-}
-
-void* Buffer::raw_ptr() {
-  if (!ptr_) {
-    return nullptr;
-  }
-  return static_cast<cu::CudaBuffer*>(ptr_)->data;
-}
-
-} // namespace allocator
-
-size_t get_active_memory() {
-  return cu::allocator().get_active_memory();
-}
-size_t get_peak_memory() {
-  return cu::allocator().get_peak_memory();
-}
-void reset_peak_memory() {
-  return cu::allocator().reset_peak_memory();
-}
-size_t set_memory_limit(size_t limit) {
-  return cu::allocator().set_memory_limit(limit);
-}
-size_t get_memory_limit() {
-  return cu::allocator().get_memory_limit();
-}
-
-// TODO: Implement buffer cache.
-size_t get_cache_memory() {
-  return 0;
-}
-size_t set_cache_limit(size_t) {
-  return 0;
-}
-size_t set_wired_limit(size_t) {
-  return 0;
-}
-void clear_cache() {}
-
-} // namespace mlx::core

mlx/backend/cuda/allocator.h

@@ -1,58 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-#include "mlx/allocator.h"
-
-#include <mutex>
-#include <set>
-#include <thread>
-#include <utility>
-
-namespace mlx::core::cu {
-
-class Worker;
-
-using allocator::Buffer;
-
-// Stores cuda-managed unified memory.
-struct CudaBuffer {
-  void* data;
-  size_t size;
-};
-
-class CudaAllocator : public allocator::Allocator {
- public:
-  Buffer malloc(size_t size) override;
-  void free(Buffer buffer) override;
-  size_t size(Buffer buffer) const override;
-
-  // Register current thread as safe to free buffers.
-  // In cuda freeing a buffer implicitly synchronizes stream, and for threads
-  // that may be waited by gpu stream (for example cpu stream threads), freeing
-  // buffers there would result in dead lock.
-  void register_this_thread();
-
-  size_t get_active_memory() const;
-  size_t get_peak_memory() const;
-  void reset_peak_memory();
-  size_t get_memory_limit();
-  size_t set_memory_limit(size_t limit);
-
- private:
-  CudaAllocator();
-  friend CudaAllocator& allocator();
-
-  std::mutex worker_mutex_;
-  std::unique_ptr<Worker> worker_;
-  std::set<std::thread::id> allowed_threads_;
-
-  std::mutex mutex_;
-  size_t memory_limit_;
-  size_t active_memory_{0};
-  size_t peak_memory_{0};
-};
-
-CudaAllocator& allocator();
-
-} // namespace mlx::core::cu

mlx/backend/cuda/copy.cpp

@@ -1,26 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/gpu/copy.h"
-
-namespace mlx::core {
-
-void copy_gpu_inplace(
-    const array& in,
-    array& out,
-    const Shape& data_shape,
-    const Strides& strides_in_pre,
-    const Strides& strides_out_pre,
-    int64_t inp_offset,
-    int64_t out_offset,
-    CopyType ctype,
-    const Stream& s,
-    const std::optional<array>& dynamic_i_offset /* = std::nullopt */,
-    const std::optional<array>& dynamic_o_offset /* = std::nullopt */) {
-  throw std::runtime_error("copy_gpu_inplace not implemented in CUDA backend.");
-}
-
-void fill_gpu(const array& val, array& out, const Stream& s) {
-  throw std::runtime_error("fill_gpu not implemented in CUDA backend.");
-}
-
-} // namespace mlx::core

mlx/backend/cuda/device.cpp

@@ -1,117 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/cuda/device.h"
-#include "mlx/backend/cuda/worker.h"
-#include "mlx/backend/metal/metal.h"
-
-#include <fmt/format.h>
-#include <nvtx3/nvtx3.hpp>
-
-namespace mlx::core {
-
-namespace cu {
-
-DeviceStream::DeviceStream(Device& device) : device_(device), stream_(device) {}
-
-void DeviceStream::synchronize() {
-  cudaStreamSynchronize(stream_);
-}
-
-cudaStream_t DeviceStream::schedule_cuda_stream() {
-  // TODO: Return a stream that maximizes parallelism.
-  return stream_;
-}
-
-cudaStream_t DeviceStream::last_cuda_stream() {
-  return stream_;
-}
-
-CommandEncoder& DeviceStream::get_encoder() {
-  if (!encoder_) {
-    encoder_ = std::make_unique<CommandEncoder>(*this);
-  }
-  return *encoder_;
-}
-
-Device::Device(int device) : device_(device) {
-  // Validate the requirements of device.
-  int attr = 0;
-  cudaDeviceGetAttribute(&attr, cudaDevAttrConcurrentManagedAccess, device_);
-  if (attr != 1) {
-    throw std::runtime_error(fmt::format(
-        "Device {} does not support synchronization in managed memory.",
-        device_));
-  }
-}
-
-void Device::make_current() {
-  // We need to set/get current CUDA device very frequently, cache it to reduce
-  // actual calls of CUDA APIs. This function assumes single-thread in host.
-  static int current = 0;
-  if (current != device_) {
-    CHECK_CUDA_ERROR(cudaSetDevice(device_));
-    current = device_;
-  }
-}
-
-DeviceStream& Device::get_stream(Stream s) {
-  auto it = streams_.find(s.index);
-  if (it == streams_.end()) {
-    it = streams_.try_emplace(s.index, *this).first;
-  }
-  return it->second;
-}
-
-CommandEncoder::CommandEncoder(DeviceStream& s)
-    : device_(s.device()), stream_(s) {}
-
-void CommandEncoder::add_completed_handler(std::function<void()> task) {
-  worker_.add_task(std::move(task));
-}
-
-void CommandEncoder::end_encoding() {
-  if (!temporaries_.empty()) {
-    add_completed_handler([temporaries = std::move(temporaries_)]() {});
-  }
-
-  // There is no kernel running, run completion handlers immediately.
-  if (!has_gpu_work_) {
-    worker_.consume_in_this_thread();
-    return;
-  }
-  has_gpu_work_ = false;
-
-  // Put completion handlers in a batch.
-  worker_.end_batch();
-
-  // Signaling kernel completion is expensive, delay until enough batches.
-  // TODO: This number is arbitrarily picked, profile for a better stragety.
-  if (worker_.uncommited_batches() > 8) {
-    commit();
-  }
-}
-
-void CommandEncoder::commit() {
-  worker_.commit(stream_.last_cuda_stream());
-}
-
-Device& device(mlx::core::Device device) {
-  static std::unordered_map<int, Device> devices;
-  auto it = devices.find(device.index);
-  if (it == devices.end()) {
-    it = devices.try_emplace(device.index, device.index).first;
-  }
-  return it->second;
-}
-
-DeviceStream& get_stream(Stream s) {
-  return device(s.device).get_stream(s);
-}
-
-CommandEncoder& get_command_encoder(Stream s) {
-  return get_stream(s).get_encoder();
-}
-
-} // namespace cu
-
-} // namespace mlx::core

mlx/backend/cuda/device.h

@@ -1,131 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-#include "mlx/array.h"
-#include "mlx/backend/cuda/worker.h"
-#include "mlx/stream.h"
-
-#include <thrust/execution_policy.h>
-
-#include <unordered_map>
-
-namespace mlx::core::cu {
-
-class Device;
-class CommandEncoder;
-
-class DeviceStream {
- public:
-  explicit DeviceStream(Device& device);
-
-  DeviceStream(const DeviceStream&) = delete;
-  DeviceStream& operator=(const DeviceStream&) = delete;
-
-  // Wait until kernels in the stream complete.
-  void synchronize();
-
-  // Return a cuda stream for launching kernels.
-  cudaStream_t schedule_cuda_stream();
-
-  // Return the last cuda stream used.
-  cudaStream_t last_cuda_stream();
-
-  CommandEncoder& get_encoder();
-
-  Device& device() {
-    return device_;
-  }
-
- private:
-  Device& device_;
-  CudaStream stream_;
-  std::unique_ptr<CommandEncoder> encoder_;
-};
-
-class Device {
- public:
-  explicit Device(int device);
-
-  Device(const Device&) = delete;
-  Device& operator=(const Device&) = delete;
-
-  // Make this device the current cuda device, required by some cuda calls.
-  void make_current();
-
-  DeviceStream& get_stream(Stream s);
-
-  int cuda_device() const {
-    return device_;
-  }
-
- private:
-  int device_;
-  std::unordered_map<int, DeviceStream> streams_;
-};
-
-class CommandEncoder {
- public:
-  explicit CommandEncoder(DeviceStream& stream);
-
-  CommandEncoder(const CommandEncoder&) = delete;
-  CommandEncoder& operator=(const CommandEncoder&) = delete;
-
-  void set_input_array(const array& arr) {}
-  void set_output_array(const array& arr) {}
-
-  void add_temporary(const array& arr) {
-    temporaries_.push_back(arr.data_shared_ptr());
-  }
-
-  void add_completed_handler(std::function<void()> task);
-  void end_encoding();
-  void commit();
-
-  // Schedule a cuda stream for |fun| to launch kernels, and check error
-  // afterwards.
-  template <typename F>
-  void launch_kernel(F&& fun) {
-    launch_kernel(stream_.schedule_cuda_stream(), std::forward<F>(fun));
-  }
-
-  template <typename F>
-  void launch_kernel(cudaStream_t stream, F&& fun) {
-    device_.make_current();
-    fun(stream);
-    check_cuda_error("kernel launch", cudaGetLastError());
-    has_gpu_work_ = true;
-  }
-
-  Device& device() {
-    return device_;
-  }
-
-  DeviceStream& stream() {
-    return stream_;
-  }
-
-  bool has_gpu_work() const {
-    return has_gpu_work_;
-  }
-
- private:
-  Device& device_;
-  DeviceStream& stream_;
-  Worker worker_;
-  bool has_gpu_work_{false};
-  std::vector<std::shared_ptr<array::Data>> temporaries_;
-};
-
-Device& device(mlx::core::Device device);
-DeviceStream& get_stream(Stream s);
-CommandEncoder& get_command_encoder(Stream s);
-
-// Return an execution policy that does not sync for result.
-// Note that not all thrust APIs support async policy, confirm before using.
-inline auto thrust_policy(cudaStream_t stream) {
-  // TODO: Connect thrust's custom allocator with mlx's allocator.
-  return thrust::cuda::par_nosync.on(stream);
-}
-
-} // namespace mlx::core::cu

mlx/backend/cuda/dtype_utils.cuh

@@ -1,35 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-#include <cuComplex.h>
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
-
-namespace mlx::core {
-
-// Maps CPU types to CUDA types.
-template <typename T>
-struct CTypeToCudaType {
-  using type = T;
-};
-
-template <>
-struct CTypeToCudaType<float16_t> {
-  using type = __half;
-};
-
-template <>
-struct CTypeToCudaType<bfloat16_t> {
-  using type = __nv_bfloat16;
-};
-
-template <>
-struct CTypeToCudaType<complex64_t> {
-  using type = cuComplex;
-};
-
-template <typename T>
-using cuda_type_t = typename CTypeToCudaType<T>::type;
-
-} // namespace mlx::core

mlx/backend/cuda/eval.cpp

@@ -1,68 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/gpu/eval.h"
-#include "mlx/backend/cuda/allocator.h"
-#include "mlx/backend/cuda/device.h"
-#include "mlx/backend/gpu/available.h"
-#include "mlx/primitives.h"
-
-#include <nvtx3/nvtx3.hpp>
-
-namespace mlx::core::gpu {
-
-bool is_available() {
-  return true;
-}
-
-void new_stream(Stream s) {
-  // Force initalization of cuda, so cuda runtime get destroyed at last.
-  cudaFree(nullptr);
-  // Ensure the static stream objects get created.
-  cu::get_command_encoder(s);
-  // The main thread is safe to free buffers.
-  cu::allocator().register_this_thread();
-}
-
-void eval(array& arr) {
-  nvtx3::scoped_range r("gpu::eval");
-  auto outputs = arr.outputs();
-  {
-    // If the array is a tracer hold a reference
-    // to its inputs so they don't get donated
-    std::vector<array> inputs;
-    if (arr.is_tracer()) {
-      inputs = arr.inputs();
-    }
-    arr.primitive().eval_gpu(arr.inputs(), outputs);
-  }
-
-  auto& encoder = cu::get_command_encoder(arr.primitive().stream());
-  if (encoder.has_gpu_work()) {
-    // 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());
-    }
-    for (auto& s : arr.siblings()) {
-      buffers.insert(s.data_shared_ptr());
-    }
-    // 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.end_encoding();
-}
-
-void finalize(Stream s) {
-  nvtx3::scoped_range r("gpu::finalize");
-  cu::get_command_encoder(s).commit();
-}
-
-void synchronize(Stream s) {
-  nvtx3::scoped_range r("gpu::synchronize");
-  cu::get_stream(s).synchronize();
-}
-
-} // namespace mlx::core::gpu

mlx/backend/cuda/event.cu

@@ -1,265 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-#include "mlx/backend/cuda/device.h"
-#include "mlx/backend/cuda/event.h"
-#include "mlx/backend/cuda/utils.h"
-#include "mlx/event.h"
-#include "mlx/scheduler.h"
-
-#include <nvtx3/nvtx3.hpp>
-
-namespace mlx::core {
-
-namespace cu {
-
-///////////////////////////////////////////////////////////////////////////////
-// CudaEvent implementations
-///////////////////////////////////////////////////////////////////////////////
-
-// Cuda event managed with RAII.
-class CudaEventHandle {
- public:
-  CudaEventHandle() {
-    CHECK_CUDA_ERROR(cudaEventCreateWithFlags(
-        &event_, cudaEventDisableTiming | cudaEventBlockingSync));
-  }
-
-  ~CudaEventHandle() {
-    CHECK_CUDA_ERROR(cudaEventDestroy(event_));
-  }
-
-  CudaEventHandle(const CudaEventHandle&) = delete;
-  CudaEventHandle& operator=(const CudaEventHandle&) = delete;
-
-  operator cudaEvent_t() const {
-    return event_;
-  }
-
- private:
-  cudaEvent_t event_;
-};
-
-CudaEvent::CudaEvent() : event_(std::make_shared<CudaEventHandle>()) {}
-
-void CudaEvent::wait() {
-  nvtx3::scoped_range r("cu::CudaEvent::wait");
-  if (!recorded_) {
-    throw std::runtime_error("Should not wait on a CudaEvent before record.");
-  }
-  cudaEventSynchronize(*event_);
-}
-
-void CudaEvent::wait(cudaStream_t stream) {
-  if (!recorded_) {
-    throw std::runtime_error("Should not wait on a CudaEvent before record.");
-  }
-  cudaStreamWaitEvent(stream, *event_);
-}
-
-void CudaEvent::wait(Stream s) {
-  if (s.device == mlx::core::Device::cpu) {
-    scheduler::enqueue(s, [*this]() mutable { wait(); });
-  } else {
-    wait(cu::get_stream(s).last_cuda_stream());
-  }
-}
-
-void CudaEvent::record(cudaStream_t stream) {
-  cudaEventRecord(*event_, stream);
-  recorded_ = true;
-}
-
-void CudaEvent::record(Stream s) {
-  if (s.device == mlx::core::Device::cpu) {
-    throw std::runtime_error("CudaEvent can not wait on cpu stream.");
-  } else {
-    record(cu::get_stream(s).last_cuda_stream());
-  }
-}
-
-bool CudaEvent::completed() const {
-  return cudaEventQuery(*event_) == cudaSuccess;
-}
-
-///////////////////////////////////////////////////////////////////////////////
-// SharedEvent implementations
-///////////////////////////////////////////////////////////////////////////////
-
-namespace {
-
-__host__ __device__ void event_wait(SharedEvent::Atomic* ac, uint64_t value) {
-  uint64_t current;
-  while ((current = ac->load()) < value) {
-    ac->wait(current);
-  }
-}
-
-__host__ __device__ void event_signal(SharedEvent::Atomic* ac, uint64_t value) {
-  ac->store(value);
-  ac->notify_all();
-}
-
-__global__ void event_wait_kernel(SharedEvent::Atomic* ac, uint64_t value) {
-  event_wait(ac, value);
-}
-
-__global__ void event_signal_kernel(SharedEvent::Atomic* ac, uint64_t value) {
-  event_signal(ac, value);
-}
-
-} // namespace
-
-SharedEvent::SharedEvent() {
-  // Allocate cuda::atomic on managed memory.
-  allocator::Buffer buffer = allocator::malloc(sizeof(Atomic));
-  Atomic* ac = static_cast<Atomic*>(buffer.raw_ptr());
-  new (ac) Atomic(0);
-  ac_ = std::shared_ptr<Atomic>(ac, [buffer](Atomic* ptr) {
-    ptr->~Atomic();
-    allocator::free(buffer);
-  });
-}
-
-void SharedEvent::wait(uint64_t value) {
-  nvtx3::scoped_range r("cu::SharedEvent::wait");
-  event_wait(ac_.get(), value);
-}
-
-void SharedEvent::wait(cudaStream_t stream, uint64_t value) {
-  event_wait_kernel<<<1, 1, 0, stream>>>(ac_.get(), value);
-}
-
-void SharedEvent::wait(Stream s, uint64_t value) {
-  nvtx3::scoped_range r("cu::SharedEvent::wait(s)");
-  if (s.device == mlx::core::Device::cpu) {
-    scheduler::enqueue(s, [*this, value]() mutable { wait(value); });
-  } else {
-    auto& encoder = get_command_encoder(s);
-    encoder.launch_kernel(
-        encoder.stream().last_cuda_stream(),
-        [this, value](cudaStream_t stream) { wait(stream, value); });
-    encoder.add_completed_handler([ac = ac_]() {});
-    encoder.end_encoding();
-  }
-}
-
-void SharedEvent::signal(uint64_t value) {
-  nvtx3::scoped_range r("cu::SharedEvent::signal");
-  event_signal(ac_.get(), value);
-}
-
-void SharedEvent::signal(cudaStream_t stream, uint64_t value) {
-  event_signal_kernel<<<1, 1, 0, stream>>>(ac_.get(), value);
-}
-
-void SharedEvent::signal(Stream s, uint64_t value) {
-  nvtx3::scoped_range r("cu::SharedEvent::signal(s)");
-  if (s.device == mlx::core::Device::cpu) {
-    scheduler::enqueue(s, [*this, value]() mutable { signal(value); });
-  } else {
-    auto& encoder = get_command_encoder(s);
-    encoder.launch_kernel(
-        encoder.stream().last_cuda_stream(),
-        [this, value](cudaStream_t stream) { signal(stream, value); });
-    encoder.add_completed_handler([ac = ac_]() {});
-    encoder.end_encoding();
-  }
-}
-
-bool SharedEvent::is_signaled(uint64_t value) const {
-  nvtx3::scoped_range r("cu::SharedEvent::is_signaled");
-  return ac_->load() >= value;
-}
-
-uint64_t SharedEvent::value() const {
-  nvtx3::scoped_range r("cu::SharedEvent::value");
-  return ac_->load();
-}
-
-} // namespace cu
-
-///////////////////////////////////////////////////////////////////////////////
-// Event implementations
-///////////////////////////////////////////////////////////////////////////////
-
-namespace {
-
-struct EventImpl {
-  // CudaEvent is preferred when possible because it is fast, however we have
-  // to fallback to SharedEvent in following cases:
-  // 1. the event is used to wait/signal a cpu stream;
-  // 2. signal value other than 1 has been specified.
-  std::unique_ptr<cu::CudaEvent> cuda;
-  std::unique_ptr<cu::SharedEvent> shared;
-
-  bool is_created() const {
-    return cuda || shared;
-  }
-
-  void ensure_created(Stream s, uint64_t signal_value) {
-    if (is_created()) {
-      return;
-    }
-    if (s.device == mlx::core::Device::cpu || signal_value > 1) {
-      nvtx3::mark("Using slow SharedEvent");
-      shared = std::make_unique<cu::SharedEvent>();
-    } else {
-      cuda = std::make_unique<cu::CudaEvent>();
-    }
-  }
-};
-
-} // namespace
-
-Event::Event(Stream s) : stream_(s) {
-  event_ = std::shared_ptr<void>(
-      new EventImpl(), [](void* ptr) { delete static_cast<EventImpl*>(ptr); });
-}
-
-void Event::wait() {
-  auto* event = static_cast<EventImpl*>(event_.get());
-  assert(event->is_created());
-  if (event->cuda) {
-    assert(value() == 1);
-    event->cuda->wait();
-  } else {
-    event->shared->wait(value());
-  }
-}
-
-void Event::wait(Stream s) {
-  auto* event = static_cast<EventImpl*>(event_.get());
-  assert(event->is_created());
-  if (event->cuda) {
-    assert(value() == 1);
-    event->cuda->wait(s);
-  } else {
-    event->shared->wait(s, value());
-  }
-}
-
-void Event::signal(Stream s) {
-  auto* event = static_cast<EventImpl*>(event_.get());
-  event->ensure_created(s, value());
-  if (event->cuda) {
-    assert(value() == 1);
-    event->cuda->record(s);
-  } else {
-    event->shared->signal(s, value());
-  }
-}
-
-bool Event::is_signaled() const {
-  auto* event = static_cast<EventImpl*>(event_.get());
-  if (!event->is_created()) {
-    return false;
-  }
-  if (event->cuda) {
-    assert(value() == 1);
-    return event->cuda->recorded() && event->cuda->completed();
-  } else {
-    return event->shared->is_signaled(value());
-  }
-}
-
-} // namespace mlx::core

mlx/backend/cuda/event.h

@@ -1,66 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-#include "mlx/stream.h"
-
-#include <cuda_runtime.h>
-#include <cuda/atomic>
-
-#include <memory>
-
-namespace mlx::core::cu {
-
-class CudaEventHandle;
-
-// Wrapper of native cuda event. It can synchronize between GPU streams, or wait
-// on GPU stream in CPU stream, but can not wait on CPU stream.
-class CudaEvent {
- public:
-  CudaEvent();
-
-  void wait();
-  void wait(cudaStream_t stream);
-  void wait(Stream s);
-  void record(cudaStream_t stream);
-  void record(Stream s);
-
-  // Return whether the recorded kernels have completed. Note that this method
-  // returns true if record() has not been called.
-  bool completed() const;
-
-  bool recorded() const {
-    return recorded_;
-  }
-
- private:
-  bool recorded_{false};
-  std::shared_ptr<CudaEventHandle> event_;
-};
-
-// Event that can synchronize between CPU and GPU. It is much slower than
-// CudaEvent so the latter should always be preferred when possible.
-class SharedEvent {
- public:
-  using Atomic = cuda::atomic<uint64_t>;
-
-  SharedEvent();
-
-  void wait(uint64_t value);
-  void wait(cudaStream_t stream, uint64_t value);
-  void wait(Stream s, uint64_t value);
-  void signal(uint64_t value);
-  void signal(cudaStream_t stream, uint64_t value);
-  void signal(Stream s, uint64_t value);
-  bool is_signaled(uint64_t value) const;
-  uint64_t value() const;
-
-  const std::shared_ptr<Atomic>& atomic() const {
-    return ac_;
-  }
-
- private:
-  std::shared_ptr<Atomic> ac_;
-};
-
-} // namespace mlx::core::cu

mlx/backend/cuda/fence.cu

@@ -1,70 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/cuda/device.h"
-#include "mlx/backend/cuda/event.h"
-#include "mlx/fence.h"
-#include "mlx/scheduler.h"
-
-#include <nvtx3/nvtx3.hpp>
-
-namespace mlx::core {
-
-namespace {
-
-__host__ __device__ void busy_wait(cuda::atomic<uint64_t>* ac, uint64_t value) {
-  while (true) {
-    // In theory the atomic_thread_fence is not needed, but for CUDA 11 without
-    // it the load() may never return new value.
-    cuda::atomic_thread_fence(cuda::memory_order_seq_cst);
-    uint64_t current = ac->load();
-    if (current >= value) {
-      break;
-    }
-  }
-}
-
-__global__ void busy_wait_kernel(cuda::atomic<uint64_t>* ac, uint64_t value) {
-  busy_wait(ac, value);
-}
-
-} // namespace
-
-struct FenceImpl {
-  uint32_t count;
-  cu::SharedEvent event;
-};
-
-Fence::Fence(Stream s) {
-  fence_ = std::shared_ptr<void>(
-      new FenceImpl{0}, [](void* ptr) { delete static_cast<FenceImpl*>(ptr); });
-}
-
-void Fence::wait(Stream s, const array&) {
-  auto* fence = static_cast<FenceImpl*>(fence_.get());
-  // We can't use SharedEvent::wait because it could hang in CUDA 11, see also:
-  // https://github.com/ml-explore/mlx/issues/2137
-  const auto& ac = fence->event.atomic();
-  if (s.device == mlx::core::Device::cpu) {
-    scheduler::enqueue(s, [ac, count = fence->count]() {
-      nvtx3::scoped_range r("Fence::wait()");
-      busy_wait(ac.get(), count);
-    });
-  } else {
-    nvtx3::scoped_range r("Fence::wait(s)");
-    auto& encoder = cu::get_command_encoder(s);
-    encoder.launch_kernel(
-        encoder.stream().last_cuda_stream(), [&](cudaStream_t stream) {
-          busy_wait_kernel<<<1, 1, 0>>>(ac.get(), fence->count);
-        });
-    encoder.add_completed_handler([ac]() {});
-    encoder.end_encoding();
-  }
-}
-
-void Fence::update(Stream s, const array&) {
-  auto* fence = static_cast<FenceImpl*>(fence_.get());
-  fence->count++;
-  fence->event.signal(s, fence->count);
-}
-
-} // namespace mlx::core

mlx/backend/cuda/kernels/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/kernels/fp16_math.cuh

@@ -1,107 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-#include <cuda_fp16.h>
-#include <cuda/std/limits>
-#include <cuda/std/type_traits>
-
-namespace mlx::core::cu {
-
-///////////////////////////////////////////////////////////////////////////////
-// Missing C++ operator overrides for CUDA 7.
-///////////////////////////////////////////////////////////////////////////////
-
-#if CUDART_VERSION < 12000 && __CUDA_ARCH__ < 800
-
-#define MLX_DEFINE_BF16_OP(OP)                                           \
-  __forceinline__ __device__ __nv_bfloat16 operator OP(                  \
-      __nv_bfloat16 x, __nv_bfloat16 y) {                                \
-    return __float2bfloat16(__bfloat162float(x) OP __bfloat162float(y)); \
-  }
-
-#define MLX_DEFINE_BF16_CMP(OP)                                          \
-  __forceinline__ __device__ bool operator OP(                           \
-      __nv_bfloat16 x, __nv_bfloat16 y) {                                \
-    return __float2bfloat16(__bfloat162float(x) OP __bfloat162float(y)); \
-  }
-
-MLX_DEFINE_BF16_OP(+)
-MLX_DEFINE_BF16_OP(-)
-MLX_DEFINE_BF16_OP(*)
-MLX_DEFINE_BF16_OP(/)
-MLX_DEFINE_BF16_CMP(>)
-MLX_DEFINE_BF16_CMP(<)
-MLX_DEFINE_BF16_CMP(>=)
-MLX_DEFINE_BF16_CMP(<=)
-
-#undef MLX_DEFINE_BF16_OP
-#undef MLX_DEFINE_BF16_CMP
-
-#endif // CUDART_VERSION < 12000 && __CUDA_ARCH__ < 800
-
-///////////////////////////////////////////////////////////////////////////////
-// Additional C++ operator overrides between half types and native types.
-///////////////////////////////////////////////////////////////////////////////
-
-template <typename T, typename U>
-constexpr bool is_integral_except =
-    cuda::std::is_integral_v<T> && !cuda::std::is_same_v<T, U>;
-
-template <typename T, typename U>
-constexpr bool is_arithmetic_except =
-    cuda::std::is_arithmetic_v<T> && !cuda::std::is_same_v<T, U>;
-
-#define MLX_DEFINE_HALF_OP(HALF, HALF2FLOAT, FLOAT2HALF, OP)          \
-  template <                                                          \
-      typename T,                                                     \
-      typename = cuda::std::enable_if_t<is_integral_except<T, HALF>>> \
-  __forceinline__ __device__ HALF operator OP(HALF x, T y) {          \
-    return FLOAT2HALF(HALF2FLOAT(x) OP static_cast<float>(y));        \
-  }                                                                   \
-  template <                                                          \
-      typename T,                                                     \
-      typename = cuda::std::enable_if_t<is_integral_except<T, HALF>>> \
-  __forceinline__ __device__ HALF operator OP(T x, HALF y) {          \
-    return FLOAT2HALF(static_cast<float>(x) OP HALF2FLOAT(y));        \
-  }
-
-#define MLX_DEFINE_HALF_CMP(HALF, HALF2FLOAT, OP)                       \
-  template <                                                            \
-      typename T,                                                       \
-      typename = cuda::std::enable_if_t<is_arithmetic_except<T, HALF>>> \
-  __forceinline__ __device__ bool operator OP(HALF x, T y) {            \
-    return HALF2FLOAT(x) OP static_cast<float>(y);                      \
-  }                                                                     \
-  template <                                                            \
-      typename T,                                                       \
-      typename = cuda::std::enable_if_t<is_arithmetic_except<T, HALF>>> \
-  __forceinline__ __device__ bool operator OP(T x, HALF y) {            \
-    return static_cast<float>(y) OP HALF2FLOAT(x);                      \
-  }
-
-MLX_DEFINE_HALF_OP(__half, __half2float, __float2half, +)
-MLX_DEFINE_HALF_OP(__half, __half2float, __float2half, -)
-MLX_DEFINE_HALF_OP(__half, __half2float, __float2half, *)
-MLX_DEFINE_HALF_OP(__half, __half2float, __float2half, /)
-MLX_DEFINE_HALF_OP(__nv_bfloat16, __bfloat162float, __float2bfloat16, +)
-MLX_DEFINE_HALF_OP(__nv_bfloat16, __bfloat162float, __float2bfloat16, -)
-MLX_DEFINE_HALF_OP(__nv_bfloat16, __bfloat162float, __float2bfloat16, *)
-MLX_DEFINE_HALF_OP(__nv_bfloat16, __bfloat162float, __float2bfloat16, /)
-MLX_DEFINE_HALF_CMP(__half, __half2float, <)
-MLX_DEFINE_HALF_CMP(__half, __half2float, >)
-MLX_DEFINE_HALF_CMP(__half, __half2float, <=)
-MLX_DEFINE_HALF_CMP(__half, __half2float, >=)
-MLX_DEFINE_HALF_CMP(__half, __half2float, ==)
-MLX_DEFINE_HALF_CMP(__half, __half2float, !=)
-MLX_DEFINE_HALF_CMP(__nv_bfloat16, __bfloat162float, <)
-MLX_DEFINE_HALF_CMP(__nv_bfloat16, __bfloat162float, >)
-MLX_DEFINE_HALF_CMP(__nv_bfloat16, __bfloat162float, <=)
-MLX_DEFINE_HALF_CMP(__nv_bfloat16, __bfloat162float, >=)
-MLX_DEFINE_HALF_CMP(__nv_bfloat16, __bfloat162float, ==)
-MLX_DEFINE_HALF_CMP(__nv_bfloat16, __bfloat162float, !=)
-
-#undef MLX_DEFINE_HALF_OP
-#undef MLX_DEFINE_HALF_CMP
-
-} // namespace mlx::core::cu

mlx/backend/cuda/primitives.cu

@@ -1,163 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/cuda/device.h"
-#include "mlx/backend/cuda/dtype_utils.cuh"
-#include "mlx/backend/cuda/kernels/arange.cuh"
-#include "mlx/backend/cuda/kernels/fp16_math.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& s = stream();
-  auto& encoder = cu::get_command_encoder(s);
-  encoder.set_output_array(out);
-  encoder.launch_kernel([&, this](cudaStream_t stream) {
-    MLX_SWITCH_INT_FLOAT_TYPES_CHECKED(out.dtype(), "Arange", CTYPE, {
-      using OutType = cuda_type_t<CTYPE>;
-      CTYPE step =
-          static_cast<CTYPE>(start_ + step_) - static_cast<CTYPE>(start_);
-      thrust::transform(
-          cu::thrust_policy(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)});
-    });
-  });
-}
-
-#define NO_GPU_MULTI(func)                                             \
-  void func::eval_gpu(                                                 \
-      const std::vector<array>& inputs, std::vector<array>& outputs) { \
-    throw std::runtime_error(#func " has no CUDA implementation.");    \
-  }
-
-#define NO_GPU(func)                                                  \
-  void func::eval_gpu(const std::vector<array>& inputs, array& out) { \
-    throw std::runtime_error(#func " has no CUDA implementation.");   \
-  }
-
-NO_GPU(Abs)
-NO_GPU(Add)
-NO_GPU(AddMM)
-NO_GPU(ArcCos)
-NO_GPU(ArcCosh)
-NO_GPU(ArcSin)
-NO_GPU(ArcSinh)
-NO_GPU(ArcTan)
-NO_GPU(ArcTan2)
-NO_GPU(ArcTanh)
-NO_GPU(ArgPartition)
-NO_GPU(ArgReduce)
-NO_GPU(ArgSort)
-NO_GPU(BitwiseBinary)
-NO_GPU(BitwiseInvert)
-NO_GPU(BlockMaskedMM)
-NO_GPU(Ceil)
-NO_GPU_MULTI(Compiled)
-NO_GPU(Conjugate)
-NO_GPU(Convolution)
-NO_GPU(Cos)
-NO_GPU(Cosh)
-NO_GPU(Divide)
-NO_GPU_MULTI(DivMod)
-NO_GPU(DynamicSlice)
-NO_GPU(DynamicSliceUpdate)
-NO_GPU(Remainder)
-NO_GPU(Equal)
-NO_GPU(Erf)
-NO_GPU(ErfInv)
-NO_GPU(Exp)
-NO_GPU(Expm1)
-NO_GPU(FFT)
-NO_GPU(Floor)
-NO_GPU(Gather)
-NO_GPU(GatherAxis)
-NO_GPU(GatherMM)
-NO_GPU(GatherQMM)
-NO_GPU(Greater)
-NO_GPU(GreaterEqual)
-NO_GPU(Hadamard)
-NO_GPU(Imag)
-NO_GPU(Less)
-NO_GPU(LessEqual)
-NO_GPU(Load)
-NO_GPU(Log)
-NO_GPU(Log1p)
-NO_GPU(LogicalNot)
-NO_GPU(LogicalAnd)
-NO_GPU(LogicalOr)
-NO_GPU(LogAddExp)
-NO_GPU(LogSumExp)
-NO_GPU_MULTI(LUF)
-NO_GPU(Matmul)
-NO_GPU(Maximum)
-NO_GPU(Minimum)
-NO_GPU(Multiply)
-NO_GPU(Negative)
-NO_GPU(NotEqual)
-NO_GPU(Partition)
-NO_GPU(Power)
-NO_GPU_MULTI(QRF)
-NO_GPU(QuantizedMatmul)
-NO_GPU(RandomBits)
-NO_GPU(Real)
-NO_GPU(Reduce)
-NO_GPU(Round)
-NO_GPU(Scan)
-NO_GPU(Scatter)
-NO_GPU(ScatterAxis)
-NO_GPU(Select)
-NO_GPU(Sigmoid)
-NO_GPU(Sign)
-NO_GPU(Sin)
-NO_GPU(Sinh)
-NO_GPU(SliceUpdate)
-NO_GPU(Softmax)
-NO_GPU(Sort)
-NO_GPU(Square)
-NO_GPU(Sqrt)
-NO_GPU(Subtract)
-NO_GPU_MULTI(SVD)
-NO_GPU(Tan)
-NO_GPU(Tanh)
-NO_GPU(Inverse)
-NO_GPU(Cholesky)
-NO_GPU_MULTI(Eigh)
-
-namespace fast {
-NO_GPU_MULTI(LayerNorm)
-NO_GPU_MULTI(LayerNormVJP)
-NO_GPU_MULTI(RMSNorm)
-NO_GPU_MULTI(RMSNormVJP)
-NO_GPU_MULTI(RoPE)
-NO_GPU(ScaledDotProductAttention)
-NO_GPU_MULTI(AffineQuantize)
-NO_GPU_MULTI(CustomKernel)
-} // namespace fast
-
-namespace distributed {
-NO_GPU_MULTI(AllReduce)
-NO_GPU_MULTI(AllGather)
-NO_GPU_MULTI(Send)
-NO_GPU_MULTI(Recv)
-} // namespace distributed
-
-} // namespace mlx::core

mlx/backend/cuda/slicing.cpp

@@ -1,15 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/gpu/slicing.h"
-
-namespace mlx::core {
-
-void concatenate_gpu(
-    const std::vector<array>& inputs,
-    array& out,
-    int axis,
-    const Stream& s) {
-  throw std::runtime_error("concatenate_gpu not implemented in CUDA backend.");
-}
-
-} // namespace mlx::core

mlx/backend/cuda/utils.cpp

@@ -1,26 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/cuda/utils.h"
-#include "mlx/backend/cuda/device.h"
-
-#include <fmt/format.h>
-
-namespace mlx::core {
-
-CudaStream::CudaStream(cu::Device& device) {
-  device.make_current();
-  CHECK_CUDA_ERROR(cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking));
-}
-
-CudaStream::~CudaStream() {
-  CHECK_CUDA_ERROR(cudaStreamDestroy(stream_));
-}
-
-void check_cuda_error(const char* name, cudaError_t err) {
-  if (err != cudaSuccess) {
-    throw std::runtime_error(
-        fmt::format("{} failed: {}", name, cudaGetErrorString(err)));
-  }
-}
-
-} // namespace mlx::core

mlx/backend/cuda/utils.h

@@ -1,36 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-#include <cuda_runtime.h>
-
-namespace mlx::core {
-
-namespace cu {
-class Device;
-}
-
-// Cuda stream managed with RAII.
-class CudaStream {
- public:
-  explicit CudaStream(cu::Device& device);
-  ~CudaStream();
-
-  CudaStream(const CudaStream&) = delete;
-  CudaStream& operator=(const CudaStream&) = delete;
-
-  operator cudaStream_t() const {
-    return stream_;
-  }
-
- private:
-  cudaStream_t stream_;
-};
-
-// Throw exception if the cuda API does not succeed.
-void check_cuda_error(const char* name, cudaError_t err);
-
-// The macro version that prints the command that failed.
-#define CHECK_CUDA_ERROR(cmd) check_cuda_error(#cmd, (cmd))
-
-} // namespace mlx::core

mlx/backend/cuda/worker.cpp

@@ -1,90 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#include "mlx/backend/cuda/worker.h"
-#include "mlx/backend/cuda/allocator.h"
-#include "mlx/backend/cuda/device.h"
-
-namespace mlx::core::cu {
-
-Worker::Worker()
-    : signal_stream_(device(mlx::core::Device::gpu)),
-      worker_(&Worker::thread_fn, this) {}
-
-Worker::~Worker() {
-  {
-    std::lock_guard lock(worker_mutex_);
-    stop_ = true;
-  }
-  worker_event_.signal(batch_ + 1);
-  worker_.join();
-}
-
-void Worker::add_task(std::function<void()> task) {
-  pending_tasks_.push_back(std::move(task));
-}
-
-void Worker::consume_in_this_thread() {
-  for (auto& task : pending_tasks_) {
-    task();
-  }
-  pending_tasks_.clear();
-}
-
-void Worker::end_batch() {
-  batch_++;
-  {
-    std::lock_guard lock(worker_mutex_);
-    worker_tasks_[batch_] = std::move(pending_tasks_);
-  }
-  uncommited_batches_++;
-}
-
-void Worker::commit() {
-  if (uncommited_batches_ == 0) {
-    return;
-  }
-  uncommited_batches_ = 0;
-  worker_event_.signal(batch_);
-}
-
-void Worker::commit(cudaStream_t stream) {
-  if (uncommited_batches_ == 0) {
-    return;
-  }
-  uncommited_batches_ = 0;
-  // Signal the |worker_event_| in |signal_stream_| after the kernels in
-  // |stream_| finish running.
-  signal_event_.record(stream);
-  signal_event_.wait(signal_stream_);
-  worker_event_.signal(signal_stream_, batch_);
-}
-
-void Worker::thread_fn() {
-  // The worker thread is safe to free buffers.
-  allocator().register_this_thread();
-
-  while (!stop_) {
-    uint64_t batch = worker_event_.value();
-    Tasks tasks;
-    {
-      std::lock_guard lock(worker_mutex_);
-      // Move tasks in signaled batches.
-      auto end = worker_tasks_.upper_bound(batch);
-      for (auto it = worker_tasks_.begin(); it != end; ++it) {
-        if (tasks.empty()) {
-          tasks = std::move(it->second);
-        } else {
-          std::move(
-              it->second.begin(), it->second.end(), std::back_inserter(tasks));
-        }
-      }
-      worker_tasks_.erase(worker_tasks_.begin(), end);
-    }
-    for (auto& task : tasks) {
-      task();
-    }
-    worker_event_.wait(batch + 1);
-  }
-}
-
-} // namespace mlx::core::cu

mlx/backend/cuda/worker.h

@@ -1,68 +0,0 @@
-// Copyright © 2025 Apple Inc.
-
-#pragma once
-
-#include "mlx/backend/cuda/event.h"
-#include "mlx/backend/cuda/utils.h"
-
-#include <functional>
-#include <map>
-#include <mutex>
-#include <thread>
-
-namespace mlx::core::cu {
-
-// Run tasks in worker thread, synchronized with cuda stream.
-class Worker {
- public:
-  Worker();
-  ~Worker();
-
-  Worker(const Worker&) = delete;
-  Worker& operator=(const Worker&) = delete;
-
-  // Add a pending |task| that will run when consumed or commited.
-  void add_task(std::function<void()> task);
-
-  // Run pending tasks immediately in current thread.
-  void consume_in_this_thread();
-
-  // Put pending tasks in a batch.
-  void end_batch();
-
-  // Inform worker thread to run current batches now.
-  void commit();
-
-  // Inform worker thread to run current batches after kernels in |stream|
-  // finish running.
-  void commit(cudaStream_t stream);
-
-  // Return how many batches have been added but not committed yet.
-  size_t uncommited_batches() const {
-    return uncommited_batches_;
-  }
-
- private:
-  void thread_fn();
-
-  uint64_t batch_{0};
-  size_t uncommited_batches_{0};
-
-  // Cuda stream and event for signaling kernel completion.
-  CudaStream signal_stream_;
-  CudaEvent signal_event_;
-
-  // Worker thread.
-  SharedEvent worker_event_;
-  std::thread worker_;
-  std::mutex worker_mutex_;
-  bool stop_{false};
-
-  // Tasks are put in |pending_tasks_| first, and then moved to
-  // |worker_tasks_| when end_batch() is called.
-  using Tasks = std::vector<std::function<void()>>;
-  Tasks pending_tasks_;
-  std::map<uint64_t, Tasks> worker_tasks_;
-};
-
-} // namespace mlx::core::cu

mlx/backend/gpu/primitives.cpp

@@ -71,12 +71,7 @@ void Contiguous::eval_gpu(const std::vector<array>& inputs, array& out) {
        (allow_col_major_ && in.flags().col_contiguous))) {
     out.copy_shared_buffer(in);
   } else {
-    out.set_data(allocator::malloc(out.nbytes()));
-    copy_gpu_inplace(
-        in,
-        out,
-        in.flags().row_contiguous ? CopyType::Vector : CopyType::General,
-        stream());
+    copy_gpu(in, out, CopyType::General);
   }
 }
 

mlx/backend/metal/fft.cpp

@@ -1,13 +1,11 @@
 // Copyright © 2024 Apple Inc.
 #include <cassert>
 #include <complex>
-#include <iostream>
 #include <map>
 #include <numeric>
 #include <set>
 
 #include "mlx/3rdparty/pocketfft.h"
-#include "mlx/backend/common/transpose.h"
 #include "mlx/backend/common/utils.h"
 #include "mlx/backend/gpu/copy.h"
 #include "mlx/backend/gpu/slicing.h"
@@ -29,7 +27,7 @@ using MTLFC = std::tuple<const void*, MTL::DataType, NS::UInteger>;
 // For strided reads/writes, coalesce at least this many complex64s
 #define MIN_COALESCE_WIDTH 4
 
-inline constexpr std::array<int, 9> supported_radices() {
+inline const std::vector<int> supported_radices() {
   // Ordered by preference in decomposition.
   return {13, 11, 8, 7, 6, 5, 4, 3, 2};
 }
@@ -51,35 +49,6 @@ std::vector<int> prime_factors(int n) {
   return factors;
 }
 
-int next_fast_n(int n) {
-  return next_power_of_2(n);
-}
-
-std::vector<int> stockham_decompose(int n) {
-  auto radices = supported_radices();
-  std::vector<int> steps(radices.size(), 0);
-  int orig_n = n;
-
-  for (int i = 0; i < radices.size(); i++) {
-    int radix = radices[i];
-
-    // Manually tuned radices for powers of 2
-    if (is_power_of_2(orig_n) && orig_n < 512 && radix > 4) {
-      continue;
-    }
-
-    while (n % radix == 0) {
-      steps[i] += 1;
-      n /= radix;
-      if (n == 1) {
-        return steps;
-      }
-    }
-  }
-
-  return {};
-}
-
 struct FourStepParams {
   bool required = false;
   bool first_step = true;
@@ -96,10 +65,9 @@ void fft_op(
     bool real,
     const FourStepParams four_step_params,
     bool inplace,
-    metal::Device& d,
     const Stream& s);
 
-struct OldFFTPlan {
+struct FFTPlan {
   int n = 0;
   // Number of steps for each radix in the Stockham decomposition
   std::vector<int> stockham;
@@ -114,104 +82,9 @@ struct OldFFTPlan {
   int n2 = 0;
 };
 
-class FFTPlan {
- public:
-  enum FFTType {
-    UNSUPPORTED,
-    NOOP,
-    STOCKHAM,
-    RADER,
-    BLUESTEIN,
-    MULTIUPLOAD_BLUESTEIN,
-    SMALL_FOUR_STEP,
-    LARGE_FOUR_STEP
-  };
-
-  FFTPlan(int n) : n_(n) {
-    // NOOP
-    if (n == 1) {
-      type_ = NOOP;
-    }
-
-    // Too large for Stockham so do four step fft for powers of 2
-    else if (n > MAX_STOCKHAM_FFT_SIZE && is_power_of_2(n)) {
-      if (n <= 1 << 20) {
-        type_ = SMALL_FOUR_STEP;
-        n2_ = n > 65536 ? 1024 : 64;
-        n1_ = n / n2_;
-        steps1_ = stockham_decompose(n1_);
-        steps2_ = stockham_decompose(n2_);
-      } else {
-        type_ = LARGE_FOUR_STEP;
-      }
-    }
-
-    // Too large and not power of 2 so do multi-upload Bluestein fft
-    else if (n > MAX_STOCKHAM_FFT_SIZE) {
-      type_ = MULTIUPLOAD_BLUESTEIN;
-      bluestein_n_ = next_fast_n(2 * n - 1);
-    }
-
-    // Stockham fft
-    else if (auto steps = stockham_decompose(n); steps.size() > 0) {
-      type_ = STOCKHAM;
-      steps_ = steps;
-    }
-
-    // Add rader but for now simply fall back to bluestein when stockham not
-    // posssible
-    else if (n > MAX_BLUESTEIN_FFT_SIZE) {
-      type_ = MULTIUPLOAD_BLUESTEIN;
-      bluestein_n_ = next_fast_n(2 * n - 1);
-    } else {
-      type_ = BLUESTEIN;
-      bluestein_n_ = next_fast_n(2 * n - 1);
-      steps_ = stockham_decompose(bluestein_n_);
-    }
-  }
-
-  FFTType type() const {
-    return type_;
-  }
-
-  int size() const {
-    return n_;
-  }
-
-  const std::vector<int>& steps() const {
-    return steps_;
-  }
-
-  int first_size() const {
-    return n1_;
-  }
-
-  const std::vector<int>& first_steps() const {
-    return steps1_;
-  }
-
-  int second_size() const {
-    return n2_;
-  }
-
-  const std::vector<int>& second_steps() const {
-    return steps2_;
-  }
-
-  int bluestein_size() const {
-    return bluestein_n_;
-  }
-
- private:
-  int n_;
-  FFTType type_;
-  std::vector<int> steps_;
-  int n1_;
-  std::vector<int> steps1_;
-  int n2_;
-  std::vector<int> steps2_;
-  int bluestein_n_;
-};
+int next_fast_n(int n) {
+  return next_power_of_2(n);
+}
 
 std::vector<int> plan_stockham_fft(int n) {
   auto radices = supported_radices();
@@ -237,12 +110,15 @@ std::vector<int> plan_stockham_fft(int n) {
   throw std::runtime_error("Unplannable");
 }
 
-OldFFTPlan plan_fft(int n) {
+FFTPlan plan_fft(int n) {
   auto radices = supported_radices();
+  std::set<int> radices_set(radices.begin(), radices.end());
 
-  OldFFTPlan plan;
+  FFTPlan plan;
   plan.n = n;
   plan.rader = std::vector<int>(radices.size(), 0);
+  auto factors = prime_factors(n);
+  int remaining_n = n;
 
   // Four Step FFT when N is too large for shared mem.
   if (n > MAX_STOCKHAM_FFT_SIZE && is_power_of_2(n)) {
@@ -252,20 +128,16 @@ OldFFTPlan plan_fft(int n) {
     plan.n2 = n > 65536 ? 1024 : 64;
     plan.n1 = n / plan.n2;
     return plan;
-  }
-
-  if (n > MAX_STOCKHAM_FFT_SIZE) {
+  } else if (n > MAX_STOCKHAM_FFT_SIZE) {
     // Otherwise we use a multi-upload Bluestein's
     plan.four_step = true;
     plan.bluestein_n = next_fast_n(2 * n - 1);
     return plan;
   }
 
-  int remaining_n = n;
-  auto factors = prime_factors(n);
   for (int factor : factors) {
     // Make sure the factor is a supported radix
-    if (std::find(radices.begin(), radices.end(), factor) == radices.end()) {
+    if (radices_set.find(factor) == radices_set.end()) {
       // We only support a single Rader factor currently
       // TODO(alexbarron) investigate weirdness with large
       // Rader sizes -- possibly a compiler issue?
@@ -282,7 +154,7 @@ OldFFTPlan plan_fft(int n) {
       for (int rf : rader_factors) {
         // We don't nest Rader's algorithm so if `factor - 1`
         // isn't Stockham decomposable we give up and do Bluestein's.
-        if (std::find(radices.begin(), radices.end(), rf) == radices.end()) {
+        if (radices_set.find(rf) == radices_set.end()) {
           plan.four_step = n > MAX_BLUESTEIN_FFT_SIZE;
           plan.bluestein_n = next_fast_n(2 * n - 1);
           plan.stockham = plan_stockham_fft(plan.bluestein_n);
@@ -300,7 +172,7 @@ OldFFTPlan plan_fft(int n) {
   return plan;
 }
 
-int compute_elems_per_thread(OldFFTPlan plan) {
+int compute_elems_per_thread(FFTPlan plan) {
   // Heuristics for selecting an efficient number
   // of threads to use for a particular mixed-radix FFT.
   auto n = plan.n;
@@ -483,11 +355,9 @@ void multi_upload_bluestein_fft(
     size_t axis,
     bool inverse,
     bool real,
-    OldFFTPlan& plan,
+    FFTPlan& plan,
     std::vector<array>& copies,
     const Stream& s) {
-  auto& d = metal::device(s.device);
-
   // TODO(alexbarron) Implement fused kernels for mutli upload bluestein's
   // algorithm
   int n = inverse ? out.shape(axis) : in.shape(axis);
@@ -550,7 +420,6 @@ void multi_upload_bluestein_fft(
       /*real=*/false,
       FourStepParams(),
       /*inplace=*/false,
-      d,
       s);
   copies.push_back(pad_temp1);
 
@@ -566,7 +435,6 @@ void multi_upload_bluestein_fft(
       /* real= */ false,
       FourStepParams(),
       /*inplace=*/true,
-      d,
       s);
 
   int offset = plan.bluestein_n - (2 * n - 1);
@@ -612,7 +480,7 @@ void four_step_fft(
     size_t axis,
     bool inverse,
     bool real,
-    OldFFTPlan& plan,
+    FFTPlan& plan,
     std::vector<array>& copies,
     const Stream& s,
     bool in_place) {
@@ -625,15 +493,7 @@ void four_step_fft(
     auto temp_shape = (real && inverse) ? out.shape() : in.shape();
     array temp(temp_shape, complex64, nullptr, {});
     fft_op(
-        in,
-        temp,
-        axis,
-        inverse,
-        real,
-        four_step_params,
-        /*inplace=*/false,
-        d,
-        s);
+        in, temp, axis, inverse, real, four_step_params, /*inplace=*/false, s);
     four_step_params.first_step = false;
     fft_op(
         temp,
@@ -643,7 +503,6 @@ void four_step_fft(
         real,
         four_step_params,
         /*inplace=*/in_place,
-        d,
         s);
     copies.push_back(temp);
   } else {
@@ -659,8 +518,9 @@ void fft_op(
     bool real,
     const FourStepParams four_step_params,
     bool inplace,
-    metal::Device& d,
     const Stream& s) {
+  auto& d = metal::device(s.device);
+
   size_t n = out.dtype() == float32 ? out.shape(axis) : in.shape(axis);
   if (n == 1) {
     out.copy_shared_buffer(in);
@@ -895,517 +755,57 @@ void fft_op(
   d.add_temporaries(std::move(copies), s.index);
 }
 
-inline int compute_elems_per_thread(int n, const std::vector<int>& steps) {
-  auto radices = supported_radices();
-  std::set<int> used_radices;
-  for (int i = 0; i < steps.size(); i++) {
-    if (steps[i] > 0) {
-      used_radices.insert(radices[i % radices.size()]);
-    }
-  }
-
-  // Manual tuning for 7/11/13
-  if (used_radices.find(7) != used_radices.end() &&
-      (used_radices.find(11) != used_radices.end() ||
-       used_radices.find(13) != used_radices.end())) {
-    return 7;
-  } else if (
-      used_radices.find(11) != used_radices.end() &&
-      used_radices.find(13) != used_radices.end()) {
-    return 11;
-  }
-
-  // TODO(alexbarron) Some really weird stuff is going on
-  // for certain `elems_per_thread` on large composite n.
-  // Possibly a compiler issue?
-  if (n == 3159)
-    return 13;
-  if (n == 3645)
-    return 5;
-  if (n == 3969)
-    return 7;
-  if (n == 1982)
-    return 5;
-
-  if (used_radices.size() == 1) {
-    return *(used_radices.begin());
-  }
-  if (used_radices.size() == 2 &&
-      (used_radices.find(11) != used_radices.end() ||
-       used_radices.find(13) != used_radices.end())) {
-    return std::accumulate(used_radices.begin(), used_radices.end(), 0) / 2;
-  }
-
-  // In all other cases use the second smallest radix.
-  return *(++used_radices.begin());
-}
-
-inline array ensure_fastest_moving_axis(
-    const array& x,
-    int axis,
-    metal::Device& d,
-    const Stream& s) {
-  // The axis is already with a stride of 1 so check that we have no overlaps
-  // and broadcasting and avoid the copy.
-  if (x.strides(axis) == 1) {
-    // This is a fairly strict test perhaps consider relaxing it in the future.
-    if (x.flags().row_contiguous || x.flags().col_contiguous) {
-      return x;
-    }
-  }
-
-  // To make it the fastest moving axis simply transpose it, then copy it and
-  // then transpose it back.
-
-  // Transpose it
-  std::vector<int> axes(x.ndim(), 0);
-  for (int ax = 0; ax < axes.size(); ax++) {
-    axes[ax] = (ax < axis) ? ax : ax + 1;
-  }
-  axes.back() = axis;
-  Shape xtshape;
-  xtshape.reserve(axes.size());
-  for (auto ax : axes) {
-    xtshape.push_back(x.shape(ax));
-  }
-  array xt(xtshape, x.dtype(), nullptr, {});
-  transpose(x, xt, axes);
-
-  // Copy it
-  array xtc(xt.shape(), x.dtype(), nullptr, {});
-  copy_gpu(
-      xt,
-      xtc,
-      xt.flags().row_contiguous ? CopyType::Vector : CopyType::General,
-      s);
-  d.add_temporary(xtc, s.index);
-
-  // Transpose it
-  for (int ax = 0; ax < axes.size(); ax++) {
-    axes[ax] = (ax < axis) ? ax : ((ax == axis) ? axes.size() - 1 : ax - 1);
-  }
-  array y(x.shape(), x.dtype(), nullptr, {});
-  transpose(xtc, y, axes);
-
-  return y;
-}
-
-inline void prepare_output_array(const array& in, array& out, int axis) {
-  // Prepare the output array such that it matches the input in terms of
-  // stride ordering. Namely we might have moved `axis` around in the `in`
-  // array. We must do the same in `out`. The difference is that we don't have
-  // to copy anything because `out` contains garbage at the moment.
-
-  if (in.flags().row_contiguous && out.flags().row_contiguous) {
-    return;
-  }
-
-  std::vector<int> axes(out.ndim(), 0);
-  for (int ax = 0; ax < axes.size(); ax++) {
-    axes[ax] = (ax < axis) ? ax : ax + 1;
-  }
-  axes.back() = axis;
-  as_transposed(out, axes);
-}
-
-void fft_stockham_inplace(
-    const FFTPlan& plan,
-    const array& in_,
-    array& out,
-    size_t axis,
-    bool inverse,
-    bool real,
-    metal::Device& d,
-    const Stream& s) {
-  // Prepare the input and output arrays such that `axis` has stride 1.
-  // Possibly copy the input but never the output as it doesn't have anything
-  // useful in it yet.
-  array in = ensure_fastest_moving_axis(in_, axis, d, s);
-  prepare_output_array(in, out, axis);
-
-  // Prepare the arguments for stockham fft
-  int n = plan.size();
-  bool power_of_2 = is_power_of_2(n);
-  int total_batch_size =
-      out.dtype() == float32 ? out.size() / n : in.size() / n;
-  auto& steps = plan.steps();
-  int elems_per_thread = compute_elems_per_thread(n, steps);
-  int threads_per_fft = ceildiv(n, elems_per_thread);
-  int tg_batch_size = std::max(MIN_THREADGROUP_MEM_SIZE / n, 1);
-  int tg_mem_size = next_power_of_2(tg_batch_size * n);
-  int batch_size = ceildiv(total_batch_size, tg_batch_size);
-  batch_size = real ? ceildiv(batch_size, 2) : batch_size; // 2 RFFTs at once
-  std::vector<MTLFC> func_consts = {
-      {&inverse, MTL::DataType::DataTypeBool, 0},
-      {&power_of_2, MTL::DataType::DataTypeBool, 1},
-      {&elems_per_thread, MTL::DataType::DataTypeInt, 2}};
-  for (int i = 0; i < steps.size(); i++) {
-    func_consts.emplace_back(&steps[i], MTL::DataType::DataTypeInt, 4 + i);
-  }
-
-  // Get the kernel
-  auto in_type = in.dtype() == float32 ? "float" : "float2";
-  auto out_type = out.dtype() == float32 ? "float" : "float2";
-  std::string hash_name;
-  std::string kname;
-  kname.reserve(64);
-  hash_name.reserve(64);
-  concatenate(kname, "fft_mem_", tg_mem_size, "_", in_type, "_", out_type);
-  concatenate(hash_name, kname, "_n", n, "_inv_", inverse);
-  auto template_def =
-      get_template_definition(kname, "fft", tg_mem_size, in_type, out_type);
-  auto kernel = get_fft_kernel(d, kname, hash_name, func_consts, template_def);
-
-  // Launch it
-  auto& compute_encoder = d.get_command_encoder(s.index);
-  compute_encoder.set_compute_pipeline_state(kernel);
-  compute_encoder.set_input_array(in, 0);
-  compute_encoder.set_output_array(out, 1);
-  compute_encoder.set_bytes(n, 2);
-  compute_encoder.set_bytes(total_batch_size, 3);
-
-  MTL::Size group_dims(1, tg_batch_size, threads_per_fft);
-  MTL::Size grid_dims(batch_size, tg_batch_size, threads_per_fft);
-  compute_encoder.dispatch_threads(grid_dims, group_dims);
-}
-
-void fft_four_step_inplace(
-    const FFTPlan& plan,
-    const array& in_,
-    array& out,
-    size_t axis,
-    bool inverse,
-    bool real,
-    metal::Device& d,
-    const Stream& s) {
-  // Prepare the input and output arrays such that `axis` has stride 1.
-  // Possibly copy the input but never the output as it doesn't have anything
-  // useful in it yet.
-  array in = ensure_fastest_moving_axis(in_, axis, d, s);
-  prepare_output_array(in, out, axis);
-
-  // Also prepare the intermediate array for the four-step fft which is
-  // implemented with 2 kernel calls.
-  array intermediate(
-      (real && inverse) ? out.shape() : in.shape(), complex64, nullptr, {});
-  intermediate.set_data(allocator::malloc(intermediate.nbytes()));
-  prepare_output_array(in, intermediate, axis);
-  d.add_temporary(intermediate, s.index);
-
-  // Make the two calls
-  for (int step = 0; step < 2; step++) {
-    // Create the parameters
-    int n1 = plan.first_size();
-    int n2 = plan.second_size();
-    int n = (step == 0) ? n1 : n2;
-    bool power_of_2 = true;
-    int total_batch_size =
-        out.dtype() == float32 ? out.size() / n : in.size() / n;
-    auto& steps = (step == 0) ? plan.first_steps() : plan.second_steps();
-    int elems_per_thread = compute_elems_per_thread(n, steps);
-    int threads_per_fft = ceildiv(n, elems_per_thread);
-    int tg_batch_size =
-        std::max(MIN_THREADGROUP_MEM_SIZE / n, MIN_COALESCE_WIDTH);
-    int tg_mem_size = next_power_of_2(tg_batch_size * n);
-    int batch_size = ceildiv(total_batch_size, tg_batch_size);
-    std::vector<MTLFC> func_consts = {
-        {&inverse, MTL::DataType::DataTypeBool, 0},
-        {&power_of_2, MTL::DataType::DataTypeBool, 1},
-        {&elems_per_thread, MTL::DataType::DataTypeInt, 2}};
-    for (int i = 0; i < steps.size(); i++) {
-      func_consts.emplace_back(&steps[i], MTL::DataType::DataTypeInt, 4 + i);
-    }
-
-    // Get the kernel
-    auto to_type = [](const array& x) {
-      return x.dtype() == float32 ? "float" : "float2";
-    };
-    auto in_type = step == 0 ? to_type(in) : to_type(intermediate);
-    auto out_type = step == 0 ? to_type(intermediate) : to_type(out);
-    std::string hash_name;
-    std::string kname;
-    kname.reserve(64);
-    hash_name.reserve(64);
-    concatenate(
-        kname,
-        "four_step_mem_",
-        tg_mem_size,
-        "_",
-        in_type,
-        "_",
-        out_type,
-        "_",
-        step,
-        (real ? "_true" : "_false"));
-    concatenate(hash_name, kname, "_n", n, "_inv_", inverse);
-    auto template_def = get_template_definition(
-        kname, "four_step_fft", tg_mem_size, in_type, out_type, step, real);
-    auto kernel =
-        get_fft_kernel(d, kname, hash_name, func_consts, template_def);
-
-    // Launch it
-    auto& compute_encoder = d.get_command_encoder(s.index);
-    compute_encoder.set_compute_pipeline_state(kernel);
-    compute_encoder.set_input_array((step == 0) ? in : intermediate, 0);
-    compute_encoder.set_output_array((step == 0) ? intermediate : out, 1);
-    compute_encoder.set_bytes(n1, 2);
-    compute_encoder.set_bytes(n2, 3);
-    compute_encoder.set_bytes(total_batch_size, 4);
-
-    MTL::Size group_dims(1, tg_batch_size, threads_per_fft);
-    MTL::Size grid_dims(batch_size, tg_batch_size, threads_per_fft);
-    compute_encoder.dispatch_threads(grid_dims, group_dims);
-  }
-}
-
-void fft_bluestein(
-    const FFTPlan& plan,
-    const array& in_,
-    array& out,
-    size_t axis,
-    bool inverse,
-    bool real,
-    metal::Device& d,
-    const Stream& s) {
-  // Prepare the input and output arrays such that `axis` has stride 1.
-  // Possibly copy the input but never the output as it doesn't have anything
-  // useful in it yet.
-  array in = ensure_fastest_moving_axis(in_, axis, d, s);
-  prepare_output_array(in, out, axis);
-
-  // Prepare the arguments for bluestein fft
-  int n = plan.bluestein_size();
-  bool power_of_2 = true;
-  int total_batch_size = out.dtype() == float32 ? out.size() / plan.size()
-                                                : in.size() / plan.size();
-  auto& steps = plan.steps();
-  int elems_per_thread = compute_elems_per_thread(n, steps);
-  int threads_per_fft = ceildiv(n, elems_per_thread);
-  int tg_batch_size = std::max(MIN_THREADGROUP_MEM_SIZE / n, 1);
-  int tg_mem_size = next_power_of_2(tg_batch_size * n);
-  int batch_size = ceildiv(total_batch_size, tg_batch_size);
-  batch_size = real ? ceildiv(batch_size, 2) : batch_size; // 2 RFFTs at once
-  std::vector<MTLFC> func_consts = {
-      {&inverse, MTL::DataType::DataTypeBool, 0},
-      {&power_of_2, MTL::DataType::DataTypeBool, 1},
-      {&elems_per_thread, MTL::DataType::DataTypeInt, 2}};
-  for (int i = 0; i < steps.size(); i++) {
-    func_consts.emplace_back(&steps[i], MTL::DataType::DataTypeInt, 4 + i);
-  }
-
-  // Get the kernel
-  auto in_type = in.dtype() == float32 ? "float" : "float2";
-  auto out_type = out.dtype() == float32 ? "float" : "float2";
-  std::string hash_name;
-  std::string kname;
-  kname.reserve(64);
-  hash_name.reserve(64);
-  concatenate(
-      kname, "bluestein_fft_mem_", tg_mem_size, "_", in_type, "_", out_type);
-  concatenate(hash_name, kname, "_n", n, "_inv_", inverse);
-  auto template_def = get_template_definition(
-      kname, "bluestein_fft", tg_mem_size, in_type, out_type);
-  auto kernel = get_fft_kernel(d, kname, hash_name, func_consts, template_def);
-
-  // Get the bluestein constants
-  auto [w_k, w_q] =
-      compute_bluestein_constants(plan.size(), plan.bluestein_size());
-  d.add_temporary(w_k, s.index);
-  d.add_temporary(w_q, s.index);
-
-  // Launch it
-  auto& compute_encoder = d.get_command_encoder(s.index);
-  compute_encoder.set_compute_pipeline_state(kernel);
-  compute_encoder.set_input_array(in, 0);
-  compute_encoder.set_output_array(out, 1);
-  compute_encoder.set_input_array(w_q, 2);
-  compute_encoder.set_input_array(w_k, 3);
-  compute_encoder.set_bytes(plan.size(), 4);
-  compute_encoder.set_bytes(n, 5);
-  compute_encoder.set_bytes(total_batch_size, 6);
-
-  MTL::Size group_dims(1, tg_batch_size, threads_per_fft);
-  MTL::Size grid_dims(batch_size, tg_batch_size, threads_per_fft);
-  compute_encoder.dispatch_threads(grid_dims, group_dims);
-}
-
-void fft_multi_upload_bluestein(
-    const FFTPlan& plan,
-    const array& in_,
-    array& out,
-    size_t axis,
-    bool inverse,
-    bool real,
-    metal::Device& d,
-    const Stream& s) {
-  // Get Bluestein's constants using the CPU (this is done in the submission
-  // thread which is pretty bad).
-  auto [w_k, w_q] =
-      compute_bluestein_constants(plan.size(), plan.bluestein_size());
-  d.add_temporary(w_k, s.index);
-  d.add_temporary(w_q, s.index);
-
-  // Prepare the input
-  auto in_shape = inverse ? out.shape() : in_.shape();
-  array in(std::move(in_shape), complex64, nullptr, {});
-  if (real && !inverse) {
-    copy_gpu(
-        in_,
-        in,
-        in_.flags().row_contiguous ? CopyType::Vector : CopyType::General,
-        s);
-    d.add_temporary(in, s.index);
-  } else if (real && inverse) {
-    int back_offset = plan.size() % 2 == 0 ? 2 : 1;
-    auto slice_shape = in.shape();
-    slice_shape[axis] -= back_offset;
-    array slice_temp(slice_shape, complex64, nullptr, {});
-    array conj_temp(in.shape(), complex64, nullptr, {});
-    Shape rstarts(in.ndim(), 0);
-    Shape rstrides(in.ndim(), 1);
-    rstarts[axis] = in.shape(axis) - back_offset;
-    rstrides[axis] = -1;
-    unary_op_gpu({in_}, conj_temp, "Conjugate", s);
-    slice_gpu(in_, slice_temp, rstarts, rstrides, s);
-    concatenate_gpu({conj_temp, slice_temp}, in, (int)axis, s);
-    d.add_temporary(conj_temp, s.index);
-  } else if (inverse) {
-    unary_op_gpu({in_}, in, "Conjugate", s);
-    d.add_temporary(in, s.index);
-  } else {
-    in.copy_shared_buffer(in_);
-  }
-
-  // Multiply with
-  Strides b_strides(in.ndim(), 0);
-  b_strides[axis] = 1;
-  array w_k_broadcast(in.shape(), complex64, nullptr, {});
-  w_k_broadcast.copy_shared_buffer(w_k, b_strides, {}, w_k.data_size());
-  array x(in.shape(), complex64, nullptr, {});
-  binary_op_gpu({in, w_k_broadcast}, x, "Multiply", s);
-  d.add_temporary(x, s.index);
-
-  // Pad
-  auto padded_shape = out.shape();
-  padded_shape[axis] = plan.bluestein_size();
-  array padded_x(padded_shape, complex64, nullptr, {});
-  auto zero = array(complex64_t{0.0f, 0.0f});
-  pad_gpu(x, zero, padded_x, {(int)axis}, {0}, s);
-  d.add_temporary(zero, s.index);
-  d.add_temporary(padded_x, s.index);
-
-  // First fft
-}
-
-void fft_op_inplace(
+void fft_op(
     const array& in,
     array& out,
     size_t axis,
     bool inverse,
     bool real,
-    metal::Device& d,
+    bool inplace,
     const Stream& s) {
-  // Get the FFT size and plan it
-  auto plan =
-      FFTPlan(out.dtype() == float32 ? out.shape(axis) : in.shape(axis));
-
-  switch (plan.type()) {
-    case FFTPlan::NOOP:
-      std::cout << "--------------> 1-size FFT <-----------------" << std::endl;
-      break;
-    case FFTPlan::STOCKHAM:
-      return fft_stockham_inplace(plan, in, out, axis, inverse, real, d, s);
-    case FFTPlan::SMALL_FOUR_STEP:
-      return fft_four_step_inplace(plan, in, out, axis, inverse, real, d, s);
-    case FFTPlan::BLUESTEIN:
-      return fft_bluestein(plan, in, out, axis, inverse, real, d, s);
-    case FFTPlan::UNSUPPORTED: {
-      std::string msg;
-      concatenate(msg, "FFT of size ", plan.size(), " not supported");
-      throw std::runtime_error(msg);
-    }
-    default:
-      std::cout << "----- NYI ----" << std::endl;
-      break;
-  }
+  fft_op(in, out, axis, inverse, real, FourStepParams(), inplace, s);
 }
 
-void nd_fft_op_inplace(
+void nd_fft_op(
     const array& in,
     array& out,
     const std::vector<size_t>& axes,
     bool inverse,
     bool real,
-    metal::Device& d,
     const Stream& s) {
-  // We are going to make and possibly reuse some intermediate arrays that will
-  // hold the intermediate fft results.
-  auto shape = inverse ? in.shape() : out.shape();
-  std::vector<array> intermediates;
-  intermediates.reserve(2);
-
-  // Utility to return either in or one of the intermediates.
-  auto get_input_array = [&](int step) -> const array& {
-    // The first step so use the input array
-    if (step == 0) {
-      return in;
-    }
-
-    return intermediates[(step - 1) % 2];
-  };
-
-  // Utility to return either out or one of the intermediates. It also informs
-  // us if we should allocate memory for that output or there is already some
-  // allocated.
-  auto get_output_array = [&](int step) -> array& {
-    // It is the final step so return the output array
-    if (step == axes.size() - 1) {
-      return out;
-    }
-
-    // We already have made an array that we can use so go ahead and use it and
-    // don't reallocate the memory.
-    if (step % 2 < intermediates.size()) {
-      return intermediates[step % 2];
-    }
-
-    array x(shape, complex64, nullptr, {});
-    x.set_data(allocator::malloc(x.nbytes()));
-    intermediates.emplace_back(std::move(x));
-    d.add_temporary(intermediates.back(), s.index);
-
-    return intermediates.back();
-  };
-
   // Perform ND FFT on GPU as a series of 1D FFTs
-  for (int step = 0; step < axes.size(); step++) {
-    auto x = get_input_array(step);
-    auto y = get_output_array(step);
-    auto step_axis = axes[inverse ? step : axes.size() - step - 1];
-    auto step_real = real && (inverse ? step == axes.size() - 1 : step == 0);
-    fft_op_inplace(x, y, step_axis, inverse, step_real, d, s);
+  auto temp_shape = inverse ? in.shape() : out.shape();
+  array temp1(temp_shape, complex64, nullptr, {});
+  array temp2(temp_shape, complex64, nullptr, {});
+  std::vector<array> temp_arrs = {temp1, temp2};
+  for (int i = axes.size() - 1; i >= 0; i--) {
+    int reverse_index = axes.size() - i - 1;
+    // For 5D and above, we don't want to reallocate our two temporary arrays
+    bool inplace = reverse_index >= 3 && i != 0;
+    // Opposite order for fft vs ifft
+    int index = inverse ? reverse_index : i;
+    size_t axis = axes[index];
+    // Mirror np.fft.(i)rfftn and perform a real transform
+    // only on the final axis.
+    bool step_real = (real && index == axes.size() - 1);
+    auto step_shape = inverse ? out.shape(axis) : in.shape(axis);
+    const array& in_arr = i == axes.size() - 1 ? in : temp_arrs[1 - i % 2];
+    array& out_arr = i == 0 ? out : temp_arrs[i % 2];
+    fft_op(in_arr, out_arr, axis, inverse, step_real, inplace, s);
   }
+
+  auto& d = metal::device(s.device);
+  d.add_temporaries(std::move(temp_arrs), s.index);
 }
 
 void FFT::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& s = stream();
-  auto& d = metal::device(s.device);
   auto& in = inputs[0];
 
-  // The FFT ops above have the *_inplace suffix. This means that the memory
-  // needs to be already allocated in the output array. Similar to
-  // copy_gpu_inplace and so on.
-  //
-  // Even though we allocate the memory, we do not necessarily want the
-  // contiguous strides so the *_inplace ops may change the strides and flags
-  // of the array but will not reallocate the memory.
-
-  out.set_data(allocator::malloc(out.nbytes()));
-
   if (axes_.size() > 1) {
-    nd_fft_op_inplace(in, out, axes_, inverse_, real_, d, s);
+    nd_fft_op(in, out, axes_, inverse_, real_, s);
   } else {
-    fft_op_inplace(in, out, axes_[0], inverse_, real_, d, s);
+    fft_op(in, out, axes_[0], inverse_, real_, /*inplace=*/false, s);
   }
 }
 

mlx/backend/metal/kernels/logsumexp.h

@@ -5,28 +5,33 @@ template <typename T, typename AccT = float, int N_READS = 4>
     const device T* in,
     device T* out,
     constant int& axis_size,
-    uint gid [[threadgroup_position_in_grid]],
-    uint _lid [[thread_position_in_threadgroup]],
+    uint2 gid [[threadgroup_position_in_grid]],
+    uint2 tid [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]],
+    uint2 _lid [[thread_position_in_threadgroup]],
     uint simd_lane_id [[thread_index_in_simdgroup]],
     uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
-  int lid = _lid;
+  int lid = _lid.x;
 
   constexpr int SIMD_SIZE = 32;
+  constexpr int elem_per_group = SIMD_SIZE * 32 * N_READS;
 
   threadgroup AccT local_max[SIMD_SIZE];
   threadgroup AccT local_normalizer[SIMD_SIZE];
 
   AccT ld[N_READS];
 
-  in += gid * size_t(axis_size) + lid * N_READS;
-  if (lid * N_READS + N_READS <= axis_size) {
+  const int axis_offset = tid.y * elem_per_group;
+  in += gid.x * size_t(axis_size) + lid * N_READS + axis_offset;
+  if (axis_offset + lid * N_READS + N_READS <= axis_size) {
     for (int i = 0; i < N_READS; i++) {
       ld[i] = AccT(in[i]);
     }
   } else {
     for (int i = 0; i < N_READS; i++) {
-      ld[i] =
-          ((lid * N_READS + i) < axis_size) ? AccT(in[i]) : Limits<AccT>::min;
+      ld[i] = ((axis_offset + lid * N_READS + i) < axis_size)
+          ? AccT(in[i])
+          : Limits<AccT>::min;
     }
   }
   if (simd_group_id == 0) {
@@ -55,6 +60,7 @@ template <typename T, typename AccT = float, int N_READS = 4>
   maxval = local_max[0];
 
   // Compute exp(x_i - maxval) and store the partial sums in local_normalizer
+  out += gid.x * grid_dim.y + tid.y;
   AccT normalizer = 0;
   for (int i = 0; i < N_READS; i++) {
     normalizer += fast::exp(ld[i] - maxval);
@@ -67,7 +73,7 @@ template <typename T, typename AccT = float, int N_READS = 4>
   if (simd_group_id == 0) {
     normalizer = simd_sum(local_normalizer[simd_lane_id]);
     if (simd_lane_id == 0) {
-      out[gid] = isinf(maxval) ? T(maxval) : T(log(normalizer) + maxval);
+      out[0] = isinf(maxval) ? T(maxval) : T(log(normalizer) + maxval);
     }
   }
 }

mlx/backend/metal/logsumexp.cpp

@@ -62,15 +62,37 @@ void LogSumExp::eval_gpu(const std::vector<array>& inputs, array& out) {
   const int n_reads = 4;
   const int looped_limit = LOGSUMEXP_LOOPED_LIMIT;
 
-  std::string kernel_name = (axis_size > looped_limit) ? "looped_" : "block_";
+  bool split = n_rows < 4 && axis_size > 4 * looped_limit;
+  bool looped = !split && axis_size > looped_limit;
+  std::string kernel_name = looped ? "looped_" : "block_";
   kernel_name += "logsumexp_";
   kernel_name += type_to_name(out);
 
   auto kernel = get_logsumexp_kernel(d, kernel_name, out);
   auto& compute_encoder = d.get_command_encoder(s.index);
+  if (split) {
+    auto tmp_size = ceildiv(axis_size, looped_limit);
+    auto tmp_shape = Shape{n_rows, static_cast<int>(tmp_size)};
+    array tmp(tmp_shape, in.dtype(), nullptr, {});
+    tmp.set_data(allocator::malloc(tmp.nbytes()));
+    size_t threadgroup_size = 1024;
+    assert(threadgroup_size <= kernel->maxTotalThreadsPerThreadgroup());
+    size_t n_threads = n_rows * threadgroup_size;
+    auto grid_dims = MTL::Size(n_threads, tmp_size, 1);
+    auto group_dims = MTL::Size(threadgroup_size, 1, 1);
+    compute_encoder.set_compute_pipeline_state(kernel);
+    compute_encoder.set_input_array(in, 0);
+    compute_encoder.set_output_array(tmp, 1);
+    compute_encoder.set_bytes(axis_size, 2);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+    d.add_temporary(tmp, s.index);
+    in = tmp;
+    axis_size = tmp_size;
+  }
+
   {
     MTL::Size grid_dims, group_dims;
-    if (axis_size <= looped_limit) {
+    if (!looped) {
       size_t threadgroup_needed = (axis_size + n_reads - 1) / n_reads;
       size_t simds_needed = (threadgroup_needed + simd_size - 1) / simd_size;
       size_t threadgroup_size = simd_size * simds_needed;

python/tests/test_ops.py

@@ -760,6 +760,10 @@ class TestOps(mlx_tests.MLXTestCase):
         x = mx.broadcast_to(mx.random.uniform(shape=(2, 1, 8)), (2, 2, 8))
         self.assertTrue(mx.allclose(mx.logsumexp(x), logsumexp(x)))
 
+        # Even larger
+        x = mx.random.uniform(shape=(4 * 4096 + 3,))
+        self.assertTrue(mx.allclose(mx.logsumexp(x), logsumexp(x)))
+
     def test_mean(self):
         x = mx.array(
             [

tests/CMakeLists.txt

@@ -9,7 +9,7 @@ FetchContent_MakeAvailable(doctest)
 
 add_executable(tests ${PROJECT_SOURCE_DIR}/tests/tests.cpp)
 
-if(MLX_BUILD_METAL OR MLX_BUILD_CUDA)
+if(MLX_BUILD_METAL)
   set(METAL_TEST_SOURCES gpu_tests.cpp)
 endif()