mlx 0.30.7

data/mlx/mlx/backend/cuda/cudnn_utils.h

@@ -0,0 +1,187 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device/config.h"
+#include "mlx/backend/cuda/utils.h"
+#include "mlx/dtype_utils.h"
+
+#include <cudnn_frontend.h>
+#include <fmt/format.h>
+
+namespace mlx::core {
+
+namespace cu {
+class CommandEncoder;
+}
+
+namespace fe = cudnn_frontend;
+
+#define CHECK_CUDNN_FE_ERROR(cmd)                                    \
+  do {                                                               \
+    auto error = cmd;                                                \
+    if (!error.is_good()) {                                          \
+      throw std::runtime_error(                                      \
+          fmt::format("{} failed: {}.", #cmd, error.get_message())); \
+    }                                                                \
+  } while (0)
+
+// Return pointer alignment of |x|'s data.
+inline uint8_t get_alignment(const array& x) {
+  uint8_t alignment = 1;
+  uintptr_t address = reinterpret_cast<uintptr_t>(gpu_ptr<void>(x));
+  for (; alignment < 32; alignment *= 2) {
+    if (address % (alignment * 2)) {
+      return alignment;
+    }
+  }
+  return alignment;
+}
+
+// Convert the type of elements in |vec| to |T|.
+template <typename T, typename Vec>
+inline std::vector<T> convert_vector(const Vec& vec) {
+  return std::vector<T>(vec.begin(), vec.end());
+}
+
+// Map dtype to cudnn data type.
+inline fe::DataType_t dtype_to_cudnn_type(Dtype dtype) {
+  switch (dtype) {
+    case int8:
+      return fe::DataType_t::INT8;
+    case int32:
+      return fe::DataType_t::INT32;
+    case uint8:
+      return fe::DataType_t::UINT8;
+    case float16:
+      return fe::DataType_t::HALF;
+    case bfloat16:
+      return fe::DataType_t::BFLOAT16;
+    case float32:
+      return fe::DataType_t::FLOAT;
+    case float64:
+      return fe::DataType_t::DOUBLE;
+    default:
+      throw std::runtime_error(
+          fmt::format(
+              "Unsupported dtype in cuDNN: {}.", dtype_to_string(dtype)));
+  }
+}
+
+// Return an array that can be used as map key for |vec| with size <= MAX_NDIM.
+//
+// There are 2 differences from the const_param util from kernel_utils.cuh:
+// 1. The rest of array is filled with 0.
+// 2. This util can be used in .cpp files.
+template <int NDIM = MAX_NDIM, typename Vec>
+inline std::array<typename Vec::value_type, NDIM> vector_key(const Vec& vec) {
+  if (vec.size() > NDIM) {
+    throw std::runtime_error(
+        fmt::format("ndim can not be larger than {}.", NDIM));
+  }
+  std::array<typename Vec::value_type, NDIM> result = {};
+  std::copy_n(vec.begin(), vec.size(), result.begin());
+  return result;
+}
+
+// Extends cuDNN graph with helpers.
+class DnnGraph : public fe::graph::Graph {
+ public:
+  DnnGraph(cudnnHandle_t handle, Dtype io_dtype, Dtype compute_dtype = float32)
+      : handle_(handle) {
+    set_io_data_type(dtype_to_cudnn_type(io_dtype));
+    set_intermediate_data_type(dtype_to_cudnn_type(compute_dtype));
+    set_compute_data_type(dtype_to_cudnn_type(compute_dtype));
+  }
+
+  // Create a cuDNN tensor description from MLX array |x|.
+  auto& tensor(
+      std::shared_ptr<fe::graph::Tensor_attributes>& attrs,
+      int64_t uid,
+      const array& x) {
+    set_tensor_attrs(attrs, uid, x);
+    return attrs;
+  }
+  auto tensor(const char* name, int64_t uid, const array& x) {
+    auto attrs = Graph::tensor(fe::graph::Tensor_attributes().set_name(name));
+    tensor(attrs, uid, x);
+    return attrs;
+  }
+
+  // Create a cuDNN tensor description from MLX array |x|, and transpose it from
+  // NHWC layout to NCHW.
+  auto& tensor_nchw(
+      std::shared_ptr<fe::graph::Tensor_attributes>& attrs,
+      int64_t uid,
+      const array& x) {
+    set_tensor_attrs_nchw(attrs, uid, x);
+    return attrs;
+  }
+  auto tensor_nchw(const char* name, int64_t uid, const array& x) {
+    auto attrs = Graph::tensor(fe::graph::Tensor_attributes().set_name(name));
+    tensor_nchw(attrs, uid, x);
+    return attrs;
+  }
+
+  // Create a 4D cuDNN tensor from 1D array, with |axis| being contiguous dim.
+  auto tensor_4d(const char* name, int64_t uid, const array& x, int axis) {
+    assert(x.ndim() == 1);
+    auto attrs = Graph::tensor(fe::graph::Tensor_attributes().set_name(name));
+    std::vector<int64_t> shape(4, 1);
+    std::vector<int64_t> strides(4, 1);
+    shape.at(axis) = x.size();
+    if (axis > 0) {
+      strides.at(axis - 1) = x.size();
+    }
+    set_tensor_attrs(attrs, uid, x, shape, strides);
+    return attrs;
+  }
+
+  // Create a cuDNN tensor for scalar.
+  auto scalar(const char* name, int64_t uid, Dtype dtype) {
+    return Graph::tensor(
+        fe::graph::Tensor_attributes()
+            .set_name(name)
+            .set_uid(uid)
+            .set_dim({1, 1, 1, 1})
+            .set_stride({1, 1, 1, 1})
+            .set_is_pass_by_value(true)
+            .set_data_type(dtype_to_cudnn_type(dtype)));
+  }
+
+  // Call this before setting notes.
+  fe::error_t prepare();
+  // Call this after setting notes.
+  fe::error_t build();
+
+  // Add cuDNN graph to CUDA graph, using native CUDA graph API.
+  fe::error_t encode_graph(
+      cu::CommandEncoder& encoder,
+      std::unordered_map<int64_t, void*> variant_pack);
+  // Add cuDNN graph to CUDA graph, using stream capture.
+  fe::error_t encode_capturing(
+      cu::CommandEncoder& encoder,
+      std::unordered_map<int64_t, void*> variant_pack);
+
+ private:
+  void* prepare_workspace(cu::CommandEncoder& encoder);
+
+  void set_tensor_attrs(
+      std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
+      int64_t uid,
+      const array& x,
+      const std::vector<int64_t>& shape,
+      const std::vector<int64_t>& strides);
+  void set_tensor_attrs(
+      std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
+      int64_t uid,
+      const array& x);
+  void set_tensor_attrs_nchw(
+      std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
+      int64_t uid,
+      const array& x);
+
+  cudnnHandle_t handle_;
+};
+
+} // namespace mlx::core

data/mlx/mlx/backend/cuda/custom_kernel.cpp

@@ -0,0 +1,379 @@
+// Copyright © 2025 Apple Inc.
+
+#include <iostream>
+
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/cuda/jit_module.h"
+#include "mlx/backend/cuda/utils.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/fast.h"
+#include "mlx/fast_primitives.h"
+
+#include <fmt/format.h>
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core::fast {
+
+namespace {
+
+constexpr const char* default_header = R"(
+#include "mlx/backend/cuda/device/utils.cuh"
+
+#include <cooperative_groups.h>
+
+#define inf cuda::std::numeric_limits<float>::infinity()
+
+)";
+
+std::string template_arguments_hash(
+    const std::vector<std::pair<std::string, TemplateArg>>& template_args) {
+  if (template_args.empty()) {
+    return "";
+  }
+
+  std::string hash;
+  hash.reserve(512);
+
+  for (const auto& [name, arg] : template_args) {
+    if (std::holds_alternative<int>(arg)) {
+      hash += fmt::format("_{}", std::get<int>(arg));
+    } else if (std::holds_alternative<bool>(arg)) {
+      hash += (std::get<bool>(arg)) ? "_t" : "_f";
+    } else if (std::holds_alternative<Dtype>(arg)) {
+      hash += "_";
+      hash += get_type_string(std::get<Dtype>(arg));
+    }
+  }
+
+  return hash;
+}
+
+std::string build_kernel(
+    const std::string& func_name,
+    const std::string& header,
+    const std::string& source,
+    const std::vector<std::string>& input_names,
+    const std::vector<array>& inputs,
+    const std::vector<std::string>& output_names,
+    const std::vector<Dtype>& output_dtypes,
+    const std::vector<std::pair<std::string, TemplateArg>>& template_args,
+    const std::vector<std::tuple<bool, bool, bool>>& shape_infos) {
+  std::string kernel_source;
+  kernel_source.reserve(header.size() + source.size() + 8192);
+  kernel_source += default_header;
+  kernel_source += header;
+  kernel_source +=
+      "namespace mlx::core::cu {\n\n"
+      "namespace cg = cooperative_groups;\n\n";
+
+  kernel_source += "__global__ void ";
+  kernel_source += func_name;
+  kernel_source += "(\n";
+
+  // Add inputs
+  for (int i = 0; i < inputs.size(); ++i) {
+    const auto& name = input_names[i];
+    const auto& arr = inputs[i];
+    kernel_source += "    const ";
+    kernel_source += dtype_to_cuda_type(arr.dtype());
+    kernel_source += "* ";
+    kernel_source += name;
+    kernel_source += ",\n";
+    // Add input shape, strides and ndim if present in the source
+    if (arr.ndim() > 0) {
+      if (std::get<0>(shape_infos[i])) {
+        kernel_source += "    const __grid_constant__ Shape ";
+        kernel_source += name;
+        kernel_source += "_shape,\n";
+      }
+      if (std::get<1>(shape_infos[i])) {
+        kernel_source += "    const __grid_constant__ Strides ";
+        kernel_source += name;
+        kernel_source += "_strides,\n";
+      }
+      if (std::get<2>(shape_infos[i])) {
+        kernel_source += "    const __grid_constant__ int ";
+        kernel_source += name;
+        kernel_source += "_ndim,\n";
+      }
+    }
+  }
+
+  // Add outputs
+  for (int i = 0; i < output_names.size(); ++i) {
+    const auto& name = output_names[i];
+    const auto& dtype = output_dtypes[i];
+    kernel_source += "    ";
+    kernel_source += dtype_to_cuda_type(dtype);
+    kernel_source += "* ";
+    kernel_source += name;
+    if (i < output_names.size() - 1) {
+      kernel_source += ",\n";
+    } else {
+      kernel_source += ") {\n";
+    }
+  }
+
+  // Set compile time constants
+  if (!template_args.empty()) {
+    for (const auto& [name, arg] : template_args) {
+      if (std::holds_alternative<int>(arg)) {
+        kernel_source +=
+            fmt::format("  constexpr int {} = {};\n", name, std::get<int>(arg));
+      } else if (std::holds_alternative<bool>(arg)) {
+        kernel_source += fmt::format(
+            "  constexpr bool {} = {};\n", name, std::get<bool>(arg));
+      } else {
+        kernel_source += fmt::format(
+            "  using {} = {};\n",
+            name,
+            dtype_to_cuda_type(std::get<Dtype>(arg)));
+      }
+    }
+    kernel_source += "\n";
+  }
+
+  kernel_source += source;
+  kernel_source += "\n}\n\n} // namespace mlx::core::cu\n";
+
+  return kernel_source;
+}
+
+} // namespace
+
+CustomKernelFunction cuda_kernel(
+    const std::string& name,
+    const std::vector<std::string>& input_names,
+    const std::vector<std::string>& output_names,
+    const std::string& source,
+    const std::string& header,
+    bool ensure_row_contiguous,
+    int shared_memory) {
+  if (output_names.empty()) {
+    throw std::invalid_argument(
+        "[custom_kernel] Must specify at least one output.");
+  }
+
+  std::vector<std::tuple<bool, bool, bool>> shape_infos;
+  for (auto& n : input_names) {
+    std::tuple<bool, bool, bool> shape_info;
+    std::get<0>(shape_info) = source.find(n + "_shape") != std::string::npos;
+    std::get<1>(shape_info) = source.find(n + "_strides") != std::string::npos;
+    std::get<2>(shape_info) = source.find(n + "_ndim") != std::string::npos;
+    shape_infos.push_back(shape_info);
+  }
+
+  return [=, shape_infos = std::move(shape_infos)](
+             const std::vector<array>& inputs,
+             const std::vector<Shape>& output_shapes,
+             const std::vector<Dtype>& output_dtypes,
+             std::tuple<int, int, int> grid,
+             std::tuple<int, int, int> threadgroup,
+             const std::vector<std::pair<std::string, TemplateArg>>&
+                 template_args = {},
+             std::optional<float> init_value = std::nullopt,
+             bool verbose = false,
+             StreamOrDevice s_ = {}) {
+    if (inputs.size() != input_names.size()) {
+      std::ostringstream msg;
+      msg << "[custom_kernel] Expected `inputs` to have size "
+          << input_names.size() << " but got size " << inputs.size() << "."
+          << std::endl;
+      throw std::invalid_argument(msg.str());
+    }
+    if (output_shapes.size() != output_names.size()) {
+      std::ostringstream msg;
+      msg << "[custom_kernel] Expected `output_shapes` to have size "
+          << output_names.size() << " but got size " << output_shapes.size()
+          << "." << std::endl;
+      throw std::invalid_argument(msg.str());
+    }
+    if (output_dtypes.size() != output_names.size()) {
+      std::ostringstream msg;
+      msg << "[custom_kernel] Expected `output_dtypes` to have size "
+          << output_names.size() << " but got size " << output_dtypes.size()
+          << "." << std::endl;
+      throw std::invalid_argument(msg.str());
+    }
+
+    auto s = to_stream(s_);
+    if (s.device != Device::gpu) {
+      throw std::invalid_argument("[custom_kernel] Only supports the GPU.");
+    }
+
+    std::string kernel_name =
+        "custom_kernel_" + name + template_arguments_hash(template_args);
+    std::string kernel_source = build_kernel(
+        kernel_name,
+        header,
+        source,
+        input_names,
+        inputs,
+        output_names,
+        output_dtypes,
+        template_args,
+        shape_infos);
+
+    if (verbose) {
+      std::cout << "Generated source code for `" << kernel_name
+                << "`:" << std::endl
+                << "```" << std::endl
+                << kernel_source << std::endl
+                << "```" << std::endl;
+    }
+
+    return array::make_arrays(
+        std::move(output_shapes),
+        std::move(output_dtypes),
+        std::make_shared<CustomKernel>(
+            s,
+            std::move(kernel_name),
+            std::move(kernel_source),
+            grid,
+            threadgroup,
+            shape_infos,
+            ensure_row_contiguous,
+            init_value,
+            std::vector<ScalarArg>{},
+            false,
+            shared_memory),
+        std::move(inputs));
+  };
+}
+
+std::vector<array> precompiled_cuda_kernel(
+    const std::string& name,
+    const std::string& compiled_source,
+    const std::vector<array>& inputs,
+    const std::vector<Shape>& output_shapes,
+    const std::vector<Dtype>& output_dtypes,
+    const std::vector<ScalarArg>& scalars,
+    std::tuple<int, int, int> grid,
+    std::tuple<int, int, int> threadgroup,
+    int shared_memory,
+    std::optional<float> init_value,
+    bool ensure_row_contiguous,
+    StreamOrDevice s) {
+  std::vector<std::tuple<bool, bool, bool>> shape_infos(
+      inputs.size(), {false, false, false});
+  return array::make_arrays(
+      output_shapes,
+      output_dtypes,
+      std::make_shared<CustomKernel>(
+          to_stream(s),
+          name,
+          compiled_source,
+          grid,
+          threadgroup,
+          shape_infos,
+          ensure_row_contiguous,
+          init_value,
+          scalars,
+          true,
+          shared_memory),
+      inputs);
+}
+
+void CustomKernel::eval_gpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  nvtx3::scoped_range r("CustomKernel::eval_gpu");
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  std::vector<array> copies;
+
+  // Allocate and initialize the output arrays
+  for (auto& out : outputs) {
+    if (init_value_) {
+      copies.emplace_back(init_value_.value(), out.dtype());
+      fill_gpu(copies.back(), out, s);
+    } else {
+      out.set_data(cu::malloc_async(out.nbytes(), encoder));
+    }
+  }
+
+  // Create the input arrays and copy if needed
+  auto check_input = [&copies, &s, this](const array& x) -> const array {
+    bool no_copy = x.flags().row_contiguous;
+    if (!ensure_row_contiguous_ || no_copy) {
+      return x;
+    } else {
+      copies.push_back(array(x.shape(), x.dtype(), nullptr, {}));
+      copy_gpu(x, copies.back(), CopyType::General, s);
+      return copies.back();
+    }
+  };
+  std::vector<array> checked_inputs;
+  for (const array& in : inputs) {
+    checked_inputs.push_back(check_input(in));
+  }
+
+  // Compile the custom kernel
+  std::string kernel_name =
+      (is_precompiled_) ? name_ : "mlx::core::cu::" + name_;
+  cu::JitModule& mod = cu::get_jit_module(
+      s.device,
+      name_,
+      [&]() {
+        return std::make_tuple(
+            is_precompiled_, source_, std::vector{kernel_name});
+      },
+      false);
+
+  // Make the arguments
+  cu::KernelArgs args;
+  for (int i = 0; i < checked_inputs.size(); i++) {
+    const array& in = checked_inputs[i];
+    auto& shape_info = shape_infos_[i];
+    args.append(in);
+    if (std::get<0>(shape_info)) {
+      args.append_ndim(in.shape());
+    }
+    if (std::get<1>(shape_info)) {
+      args.append_ndim(in.strides());
+    }
+    if (std::get<2>(shape_info)) {
+      args.append<int32_t>(in.ndim());
+    }
+  }
+  for (auto& out : outputs) {
+    args.append(out);
+  }
+  for (auto& s : scalar_arguments_) {
+    if (std::holds_alternative<bool>(s)) {
+      args.append(std::get<bool>(s));
+    } else if (std::holds_alternative<int>(s)) {
+      args.append(std::get<int>(s));
+    } else if (std::holds_alternative<float>(s)) {
+      args.append(std::get<float>(s));
+    }
+  }
+
+  // Make the grid
+  const auto [tx, ty, tz] = threadgroup_;
+  const auto [gx, gy, gz] = grid_;
+  dim3 block(std::min(tx, gx), std::min(ty, gy), std::min(tz, gz));
+  dim3 grid((gx + tx - 1) / tx, (gy + ty - 1) / ty, (gz + tz - 1) / tz);
+
+  // Call the kernel
+  for (const auto& in : checked_inputs) {
+    encoder.set_input_array(in);
+  }
+  for (const auto& out : outputs) {
+    encoder.set_output_array(out);
+  }
+  for (const auto& t : copies) {
+    encoder.add_temporary(t);
+  }
+  auto kernel =
+      mod.get_kernel(kernel_name, [smem = shared_memory_](CUfunction kernel) {
+        if (smem > 0 && smem > 48000) {
+          cuFuncSetAttribute(
+              kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem);
+        }
+      });
+  encoder.add_kernel_node(kernel, grid, block, shared_memory_, args.args());
+}
+
+} // namespace mlx::core::fast

data/mlx/mlx/backend/cuda/cutlass_utils.cuh

@@ -0,0 +1,46 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/dtype.h"
+
+#include <cutlass/bfloat16.h>
+#include <cutlass/half.h>
+#include <fmt/format.h>
+
+namespace mlx::core {
+
+// Throw exception if the cutlass API does not succeed.
+inline void check_cutlass_error(const char* name, cutlass::Status status) {
+  if (status != cutlass::Status::kSuccess) {
+    throw std::runtime_error(
+        fmt::format(
+            "{} failed with code: {}.",
+            name,
+            cutlass::cutlassGetStatusString(status)));
+  }
+}
+
+// The macro version that prints the command that failed.
+#define CHECK_CUTLASS_ERROR(cmd) check_cutlass_error(#cmd, (cmd))
+
+// Maps CPU types to CUTLASS types.
+template <typename T>
+struct CTypeToCutlassType {
+  using type = T;
+};
+
+template <>
+struct CTypeToCutlassType<float16_t> {
+  using type = cutlass::half_t;
+};
+
+template <>
+struct CTypeToCutlassType<bfloat16_t> {
+  using type = cutlass::bfloat16_t;
+};
+
+template <typename T>
+using cutlass_type_t = typename CTypeToCutlassType<T>::type;
+
+} // namespace mlx::core

data/mlx/mlx/backend/cuda/delayload.cpp

@@ -0,0 +1,80 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/common/utils.h"
+
+// clang-format off
+#include <windows.h> // must be included first
+#include <delayimp.h>
+// clang-format on
+
+namespace mlx::core {
+
+namespace fs = std::filesystem;
+
+inline fs::path relative_to_current_binary(const char* relative) {
+  return fs::absolute(current_binary_dir() / relative);
+}
+
+inline fs::path cublas_bin_dir() {
+#if defined(MLX_CUDA_BIN_DIR)
+  return MLX_CUDA_BIN_DIR;
+#else
+  return relative_to_current_binary("../nvidia/cublas/bin");
+#endif
+}
+
+fs::path load_nvrtc() {
+#if defined(MLX_CUDA_BIN_DIR)
+  fs::path nvrtc_bin_dir = MLX_CUDA_BIN_DIR;
+#else
+  fs::path nvrtc_bin_dir =
+      relative_to_current_binary("../nvidia/cuda_nvrtc/bin");
+#endif
+  // Internally nvrtc loads some libs dynamically, add to search dirs.
+  ::AddDllDirectory(nvrtc_bin_dir.c_str());
+  return nvrtc_bin_dir;
+}
+
+fs::path load_cudnn() {
+#if defined(MLX_CUDNN_BIN_DIR)
+  fs::path cudnn_bin_dir = MLX_CUDNN_BIN_DIR;
+#else
+  fs::path cudnn_bin_dir = relative_to_current_binary("../nvidia/cudnn/bin");
+#endif
+  // Must load cudnn_graph64_9.dll before locating symbols, otherwise We would
+  // get errors like "Invalid handle. Cannot load symbol cudnnCreate".
+  for (const auto& dll : fs::directory_iterator(cudnn_bin_dir)) {
+    if (dll.path().filename().string().starts_with("cudnn_graph") &&
+        dll.path().extension() == ".dll") {
+      ::LoadLibraryW(dll.path().c_str());
+      break;
+    }
+  }
+  // Internally cuDNN loads some libs dynamically, add to search dirs.
+  load_nvrtc();
+  ::AddDllDirectory(cudnn_bin_dir.c_str());
+  ::AddDllDirectory(cublas_bin_dir().c_str());
+  return cudnn_bin_dir;
+}
+
+// Called by system when failed to locate a lazy-loaded DLL.
+FARPROC WINAPI delayload_helper(unsigned dliNotify, PDelayLoadInfo pdli) {
+  HMODULE mod = NULL;
+  if (dliNotify == dliNotePreLoadLibrary) {
+    std::string dll = pdli->szDll;
+    if (dll.starts_with("cudnn")) {
+      static auto cudnn_bin_dir = load_cudnn();
+      mod = ::LoadLibraryW((cudnn_bin_dir / dll).c_str());
+    } else if (dll.starts_with("cublas")) {
+      mod = ::LoadLibraryW((cublas_bin_dir() / dll).c_str());
+    } else if (dll.starts_with("nvrtc")) {
+      static auto nvrtc_bin_dir = load_nvrtc();
+      mod = ::LoadLibraryW((nvrtc_bin_dir / dll).c_str());
+    }
+  }
+  return reinterpret_cast<FARPROC>(mod);
+}
+
+} // namespace mlx::core
+
+extern "C" const PfnDliHook __pfnDliNotifyHook2 = mlx::core::delayload_helper;