mlx 1.0.0

data/mlx/mlx/backend/cuda/copy.cu

@@ -0,0 +1,132 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/cuda/copy/copy.cuh"
+
+namespace mlx::core {
+
+void copy_gpu(const array& in, array& out, CopyType ctype, const Stream& s) {
+  auto& encoder = cu::get_command_encoder(s);
+  bool donated = set_copy_output_data(
+      in, out, ctype, [&](auto n) { return cu::malloc_async(n, encoder); });
+  if (donated && in.dtype() == out.dtype()) {
+    // If the output has the same type as the input then there is nothing to
+    // copy, just use the buffer.
+    return;
+  }
+  if (ctype == CopyType::GeneralGeneral) {
+    ctype = CopyType::General;
+  }
+  copy_gpu_inplace(in, out, ctype, s);
+}
+
+void copy_gpu_inplace(
+    const array& in,
+    array& out,
+    const Shape& shape,
+    const Strides& strides_in,
+    const Strides& strides_out,
+    int64_t offset_in,
+    int64_t offset_out,
+    CopyType ctype,
+    const Stream& s,
+    std::optional<array> dynamic_offset_in,
+    std::optional<array> dynamic_offset_out) {
+  if (out.size() == 0) {
+    return;
+  }
+
+  auto& encoder = cu::get_command_encoder(s);
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  if (ctype == CopyType::Scalar || ctype == CopyType::Vector) {
+    copy_contiguous(encoder, ctype, in, out, offset_in, offset_out);
+    return;
+  }
+
+  if (ctype == CopyType::General || ctype == CopyType::GeneralGeneral) {
+    auto [shape_collapsed, strides_vec] = collapse_contiguous_dims(
+        shape, std::vector{strides_in, strides_out}, INT32_MAX);
+    if (ctype == CopyType::General) {
+      copy_general_input(
+          encoder,
+          ctype,
+          in,
+          out,
+          offset_in,
+          offset_out,
+          shape_collapsed,
+          strides_vec[0]);
+    } else {
+      if (dynamic_offset_in || dynamic_offset_out) {
+        if (!dynamic_offset_in) {
+          dynamic_offset_in = array(0, int64);
+          encoder.add_temporary(*dynamic_offset_in);
+        }
+        if (!dynamic_offset_out) {
+          dynamic_offset_out = array(0, int64);
+          encoder.add_temporary(*dynamic_offset_out);
+        }
+        encoder.set_input_array(*dynamic_offset_in);
+        encoder.set_input_array(*dynamic_offset_out);
+        copy_general_dynamic(
+            encoder,
+            ctype,
+            in,
+            out,
+            offset_in,
+            offset_out,
+            shape_collapsed,
+            strides_vec[0],
+            strides_vec[1],
+            *dynamic_offset_in,
+            *dynamic_offset_out);
+      } else {
+        copy_general(
+            encoder,
+            ctype,
+            in,
+            out,
+            offset_in,
+            offset_out,
+            shape_collapsed,
+            strides_vec[0],
+            strides_vec[1]);
+      }
+    }
+    return;
+  }
+}
+
+void fill_gpu(const array& in, array& out, const Stream& s) {
+  if (out.size() == 0) {
+    return;
+  }
+  auto& encoder = cu::get_command_encoder(s);
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+  encoder.set_input_array(in);
+  encoder.set_output_array(out);
+  copy_contiguous(encoder, CopyType::Scalar, in, out, 0, 0);
+}
+
+void reshape_gpu(const array& in, array& out, Stream s) {
+  auto [copy_necessary, out_strides] = prepare_reshape(in, out);
+  if (copy_necessary) {
+    auto& encoder = cu::get_command_encoder(s);
+    out.set_data(cu::malloc_async(out.nbytes(), encoder));
+    copy_gpu_inplace(
+        in,
+        out,
+        in.shape(),
+        in.strides(),
+        make_contiguous_strides(in.shape()),
+        0,
+        0,
+        CopyType::General,
+        s);
+  } else {
+    shared_buffer_reshape(in, out_strides, out);
+  }
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cuda/cublas_utils.cpp

@@ -0,0 +1,222 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/cublas_utils.h"
+#include "mlx/backend/cuda/cuda.h"
+#include "mlx/utils.h"
+
+namespace mlx::core {
+namespace cublas_utils {
+
+namespace {
+
+struct CublasPreference {
+  CublasPreference(cu::Device& device) {
+    // The recommended cublas workspace size is 4 MiB for pre-Hopper and 32 MiB
+    // for Hopper+:
+    // https://docs.nvidia.com/cuda/cublas/#cublassetworkspace
+    uint64_t MiB = 1024 * 1024;
+    uint64_t workspace_size =
+        device.compute_capability_major() >= 9 ? 32 * MiB : 4 * MiB;
+
+    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceCreate(&pref_));
+    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceSetAttribute(
+        pref_,
+        CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
+        &workspace_size,
+        sizeof(uint64_t)));
+  }
+
+  ~CublasPreference() {
+    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceDestroy(pref_));
+  }
+
+  cublasLtMatmulPreference_t pref_{nullptr};
+};
+
+} // namespace
+
+cublasLtMatmulPreference_t get_preference(cu::Device& device) {
+  static CublasPreference pref(device);
+  return pref.pref_;
+}
+
+cublasLtMatrixLayout_t create_matrix_layout(
+    cudaDataType_t type,
+    uint64_t rows,
+    uint64_t cols,
+    bool transposed,
+    int64_t ld,
+    int32_t batch_count,
+    int64_t batch_stride) {
+  cublasLtMatrixLayout_t desc;
+  if (transposed) {
+    std::swap(rows, cols);
+  }
+  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutCreate(&desc, type, rows, cols, ld));
+  if (batch_count > 1) {
+    CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutSetAttribute(
+        desc,
+        CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT,
+        &batch_count,
+        sizeof(int32_t)));
+    CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutSetAttribute(
+        desc,
+        CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET,
+        &batch_stride,
+        sizeof(int64_t)));
+  }
+  return desc;
+}
+
+} // namespace cublas_utils
+
+CublasMatmulBase::~CublasMatmulBase() {
+  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(a_desc_));
+  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(b_desc_));
+  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(c_desc_));
+  CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(out_desc_));
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescDestroy(matmul_desc_));
+}
+
+void CublasMatmulBase::init_base(
+    cu::Device& device,
+    cudaDataType_t scale_type,
+    cublasComputeType_t compute_type,
+    cudaDataType_t data_type,
+    cudaDataType_t output_type,
+    bool a_transposed,
+    uint64_t a_rows,
+    uint64_t a_cols,
+    int64_t lda,
+    bool b_transposed,
+    uint64_t b_rows,
+    uint64_t b_cols,
+    int64_t ldb,
+    int32_t batch_count,
+    int64_t a_batch_stride,
+    int64_t b_batch_stride) {
+  M_ = a_rows;
+  N_ = b_cols;
+  scale_type_ = scale_type;
+  handle_ = device.get_cublaslt_handle();
+  pref_ = cublas_utils::get_preference(device);
+  heuristic_.state = CUBLAS_STATUS_NOT_INITIALIZED;
+
+  CHECK_CUBLAS_ERROR(
+      cublasLtMatmulDescCreate(&matmul_desc_, compute_type, scale_type));
+
+  int32_t pointer_mode = CUBLASLT_POINTER_MODE_HOST;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_POINTER_MODE,
+      &pointer_mode,
+      sizeof(int32_t)));
+
+  // In cublasLt matrices use column-major layout, while it is possible to use
+  // the CUBLASLT_ORDER_ROW option to switch to row-major layout, the bias
+  // epilogue does not work with the option. So instead we swap A and B to make
+  // cublasLt return the row-major result, which works because:
+  // - the data of a matrix in row-major layout is identical to its transpose in
+  //   column-major layout
+  // - C^T = (A @ B)^T = B^T @ A^T
+  cublasOperation_t a_op = b_transposed ? CUBLAS_OP_T : CUBLAS_OP_N;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_TRANSA,
+      &a_op,
+      sizeof(cublasOperation_t)));
+  cublasOperation_t b_op = a_transposed ? CUBLAS_OP_T : CUBLAS_OP_N;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_TRANSB,
+      &b_op,
+      sizeof(cublasOperation_t)));
+
+  a_desc_ = cublas_utils::create_matrix_layout(
+      data_type,
+      b_cols,
+      b_rows,
+      b_transposed,
+      ldb,
+      batch_count,
+      b_batch_stride);
+  b_desc_ = cublas_utils::create_matrix_layout(
+      data_type,
+      a_cols,
+      a_rows,
+      a_transposed,
+      lda,
+      batch_count,
+      a_batch_stride);
+  out_desc_ = cublas_utils::create_matrix_layout(
+      output_type, b_cols, a_rows, false, b_cols, batch_count, b_cols * a_rows);
+}
+
+void CublasMatmulBase::execute_matmul(
+    cu::CommandEncoder& encoder,
+    void* out,
+    const void* a,
+    const void* b,
+    const void* c,
+    const void* alpha_ptr,
+    const void* beta_ptr) {
+  if (heuristic_.state != CUBLAS_STATUS_SUCCESS) {
+    int ret = 0;
+    CHECK_CUBLAS_ERROR(cublasLtMatmulAlgoGetHeuristic(
+        handle_,
+        matmul_desc_,
+        a_desc_,
+        b_desc_,
+        c ? c_desc_ : out_desc_,
+        out_desc_,
+        pref_,
+        1,
+        &heuristic_,
+        &ret));
+    if (ret == 0) {
+      throw std::runtime_error("Can not find algorithm for matmul.");
+    }
+  }
+
+  void* workspace_ptr = allocate_workspace(encoder, heuristic_.workspaceSize);
+
+  // Execute matmul
+  auto capture = encoder.capture_context();
+  CHECK_CUBLAS_ERROR(cublasLtMatmul(
+      handle_,
+      matmul_desc_,
+      alpha_ptr,
+      b, // a and b are swapped for row-major layout
+      a_desc_,
+      a,
+      b_desc_,
+      beta_ptr,
+      c ? c : out,
+      c ? c_desc_ : out_desc_,
+      out,
+      out_desc_,
+      &heuristic_.algo,
+      workspace_ptr,
+      heuristic_.workspaceSize,
+      encoder.stream()));
+}
+
+void CublasMatmulBase::set_bias(
+    cu::CommandEncoder& encoder,
+    const array& bias) {
+  encoder.set_input_array(bias);
+  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_BIAS;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_EPILOGUE,
+      &epilogue,
+      sizeof(epilogue)));
+  auto* bias_ptr = gpu_ptr<void>(bias);
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_BIAS_POINTER,
+      &bias_ptr,
+      sizeof(bias_ptr)));
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cuda/cublas_utils.h

@@ -0,0 +1,95 @@
+// Copyright © 2025 Apple Inc.
+#pragma once
+
+#include <cublasLt.h>
+#include "mlx/array.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/dtype_utils.h"
+
+namespace mlx::core {
+namespace cublas_utils {
+
+// Get the shared cublas preference for a device
+cublasLtMatmulPreference_t get_preference(cu::Device& device);
+
+cublasLtMatrixLayout_t create_matrix_layout(
+    cudaDataType_t type,
+    uint64_t rows,
+    uint64_t cols,
+    bool transposed,
+    int64_t ld,
+    int32_t batch_count,
+    int64_t batch_stride);
+
+inline cudaDataType_t dtype_to_cublas_type(Dtype dtype, std::string_view tag) {
+  switch (dtype) {
+    case float16:
+      return CUDA_R_16F;
+    case bfloat16:
+      return CUDA_R_16BF;
+    case float32:
+      return CUDA_R_32F;
+    case float64:
+      return CUDA_R_64F;
+    case complex64:
+      return CUDA_C_32F;
+    default:
+      throw std::runtime_error(
+          fmt::format(
+              "Unsupported dtype in {}: {}.", tag, dtype_to_string(dtype)));
+  }
+}
+
+} // namespace cublas_utils
+
+class CublasMatmulBase {
+ public:
+  virtual ~CublasMatmulBase();
+
+  void set_bias(cu::CommandEncoder& encoder, const array& bias);
+
+ protected:
+  CublasMatmulBase() = default;
+
+  // Common member variables shared by all matmul types
+  uint64_t M_;
+  uint64_t N_;
+  cudaDataType_t scale_type_;
+  cublasLtMatmulPreference_t pref_{nullptr};
+  cublasLtHandle_t handle_{nullptr};
+  cublasLtMatmulDesc_t matmul_desc_{nullptr};
+  cublasLtMatrixLayout_t a_desc_{nullptr};
+  cublasLtMatrixLayout_t b_desc_{nullptr};
+  cublasLtMatrixLayout_t c_desc_{nullptr};
+  cublasLtMatrixLayout_t out_desc_{nullptr};
+  cublasLtMatmulHeuristicResult_t heuristic_;
+
+  void init_base(
+      cu::Device& device,
+      cudaDataType_t scale_type,
+      cublasComputeType_t compute_type,
+      cudaDataType_t data_type,
+      cudaDataType_t output_type,
+      bool a_transposed,
+      uint64_t a_rows,
+      uint64_t a_cols,
+      int64_t lda,
+      bool b_transposed,
+      uint64_t b_rows,
+      uint64_t b_cols,
+      int64_t ldb,
+      int32_t batch_count,
+      int64_t a_batch_stride,
+      int64_t b_batch_stride);
+
+  void execute_matmul(
+      cu::CommandEncoder& encoder,
+      void* out,
+      const void* a,
+      const void* b,
+      const void* c,
+      const void* alpha_ptr,
+      const void* beta_ptr);
+};
+
+} // namespace mlx::core

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

@@ -0,0 +1,21 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <string>
+#include <unordered_map>
+#include <variant>
+
+#include "mlx/api.h"
+
+namespace mlx::core::cu {
+
+/* Check if the CUDA backend is available. */
+MLX_API bool is_available();
+
+/* Get information about a CUDA device. */
+MLX_API const
+    std::unordered_map<std::string, std::variant<std::string, size_t>>&
+    device_info(int device_index = 0);
+
+} // namespace mlx::core::cu

data/mlx/mlx/backend/cuda/cuda_utils.h

@@ -0,0 +1,90 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include <cublasLt.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cudnn.h>
+
+namespace mlx::core {
+
+// Throw exception if the cuda API does not succeed.
+void check_cublas_error(const char* name, cublasStatus_t err);
+void check_cuda_error(const char* name, cudaError_t err);
+void check_cuda_error(const char* name, CUresult err);
+void check_cudnn_error(const char* name, cudnnStatus_t err);
+
+// The macro version that prints the command that failed.
+#define CHECK_CUBLAS_ERROR(cmd) check_cublas_error(#cmd, (cmd))
+#define CHECK_CUDA_ERROR(cmd) check_cuda_error(#cmd, (cmd))
+#define CHECK_CUDNN_ERROR(cmd) check_cudnn_error(#cmd, (cmd))
+
+// Base class for RAII managed CUDA resources.
+template <typename Handle, cudaError_t (*Destroy)(Handle)>
+class CudaHandle {
+ public:
+  CudaHandle(Handle handle = nullptr) : handle_(handle) {}
+
+  CudaHandle(CudaHandle&& other) : handle_(other.handle_) {
+    assert(this != &other);
+    other.handle_ = nullptr;
+  }
+
+  ~CudaHandle() {
+    // Skip if there was an error to avoid throwing in the destructors
+    if (cudaPeekAtLastError() != cudaSuccess) {
+      return;
+    }
+    reset();
+  }
+
+  CudaHandle(const CudaHandle&) = delete;
+  CudaHandle& operator=(const CudaHandle&) = delete;
+
+  CudaHandle& operator=(CudaHandle&& other) {
+    assert(this != &other);
+    reset();
+    std::swap(handle_, other.handle_);
+    return *this;
+  }
+
+  void reset() {
+    if (handle_ != nullptr) {
+      CHECK_CUDA_ERROR(Destroy(handle_));
+      handle_ = nullptr;
+    }
+  }
+
+  operator Handle() const {
+    return handle_;
+  }
+
+ protected:
+  Handle handle_;
+};
+
+namespace cu {
+class Device;
+}; // namespace cu
+
+// Wrappers of CUDA resources.
+class CudaGraph : public CudaHandle<cudaGraph_t, cudaGraphDestroy> {
+ public:
+  using CudaHandle::CudaHandle;
+  explicit CudaGraph(cu::Device& device);
+  void end_capture(cudaStream_t stream);
+};
+
+class CudaGraphExec : public CudaHandle<cudaGraphExec_t, cudaGraphExecDestroy> {
+ public:
+  void instantiate(cudaGraph_t graph);
+};
+
+class CudaStream : public CudaHandle<cudaStream_t, cudaStreamDestroy> {
+ public:
+  using CudaHandle::CudaHandle;
+  explicit CudaStream(cu::Device& device);
+};
+
+} // namespace mlx::core

data/mlx/mlx/backend/cuda/cudnn_utils.cpp

@@ -0,0 +1,133 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/cudnn_utils.h"
+#include "mlx/backend/cuda/device.h"
+
+namespace mlx::core {
+
+namespace {
+
+#define RETURN_IF_ERROR(cmd)          \
+  if (auto ret = cmd; ret.is_bad()) { \
+    return ret;                       \
+  }
+
+// In MLX a singleton dim (shape[dim] == 1) can have any stride, but in cuDNN
+// whether a tensor is contiguous is determined with:
+// shape[dim] == shape[dim + 1] * strides[dim + 1]
+// So a contiguous array with singleton dims in MLX may be mistakenly treated
+// as strided in cuDNN, and we work around it by normalizing the strides.
+std::vector<int64_t> normalized_strides(const array& x) {
+  std::vector<int64_t> strides(x.strides().begin(), x.strides().end());
+  if (std::all_of(
+          strides.begin(), strides.end(), [](int64_t s) { return s == 0; })) {
+    strides.back() = 1;
+    return strides;
+  }
+  if (!x.flags().row_contiguous || x.ndim() < 2) {
+    return strides;
+  }
+  for (int i = x.ndim() - 2; i >= 0; --i) {
+    if (x.shape(i) == 1) {
+      strides[i] = x.shape(i + 1) * strides[i + 1];
+    }
+  }
+  return strides;
+}
+
+// Return the shape and strides after transposing from NHWC to NCHW.
+inline auto nhwc_to_nchw(const array& x) {
+  auto shape = convert_vector<int64_t>(x.shape());
+  auto strides = normalized_strides(x);
+  assert(shape.size() >= 3);
+  shape.insert(shape.begin() + 1, shape.back());
+  shape.erase(shape.end() - 1);
+  strides.insert(strides.begin() + 1, strides.back());
+  strides.erase(strides.end() - 1);
+  return std::make_tuple(std::move(shape), std::move(strides));
+}
+
+} // namespace
+
+fe::error_t DnnGraph::prepare() {
+  RETURN_IF_ERROR(validate());
+  try {
+    RETURN_IF_ERROR(build_operation_graph(handle_));
+  } catch (cudnn_frontend::cudnnException& error) {
+    // cuDNN bug: they did not catch all exceptions in the API.
+    return {fe::error_code_t::CUDNN_BACKEND_API_FAILED, error.what()};
+  }
+  RETURN_IF_ERROR(create_execution_plans({fe::HeurMode_t::A}));
+  return {};
+}
+
+fe::error_t DnnGraph::build() {
+  RETURN_IF_ERROR(check_support(handle_));
+  RETURN_IF_ERROR(build_plans(handle_));
+  return {};
+}
+
+fe::error_t DnnGraph::encode_graph(
+    cu::CommandEncoder& encoder,
+    std::unordered_map<int64_t, void*> variant_pack) {
+  cudnnSetStream(handle_, encoder.stream());
+  CudaGraph cuda_graph(encoder.device());
+  RETURN_IF_ERROR(populate_cuda_graph(
+      handle_, variant_pack, prepare_workspace(encoder), cuda_graph));
+  encoder.add_graph_node(cuda_graph);
+  return {};
+}
+
+fe::error_t DnnGraph::encode_capturing(
+    cu::CommandEncoder& encoder,
+    std::unordered_map<int64_t, void*> variant_pack) {
+  auto* workspace_ptr = prepare_workspace(encoder);
+  auto capture = encoder.capture_context();
+  cudnnSetStream(handle_, encoder.stream());
+  auto ret = execute(handle_, variant_pack, workspace_ptr);
+  if (ret.is_bad()) {
+    capture.discard = true;
+  }
+  return ret;
+}
+
+void* DnnGraph::prepare_workspace(cu::CommandEncoder& encoder) {
+  int64_t workspace_size = 0;
+  CHECK_CUDNN_FE_ERROR(get_workspace_size(workspace_size));
+  return allocate_workspace(encoder, workspace_size);
+}
+
+void DnnGraph::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) {
+  tensor->set_uid(uid)
+      .set_alignment(get_alignment(x))
+      .set_data_type(dtype_to_cudnn_type(x.dtype()))
+      .set_dim(shape)
+      .set_stride(strides);
+}
+
+void DnnGraph::set_tensor_attrs(
+    std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
+    int64_t uid,
+    const array& x) {
+  set_tensor_attrs(
+      tensor,
+      uid,
+      x,
+      convert_vector<int64_t>(x.shape()),
+      normalized_strides(x));
+}
+
+void DnnGraph::set_tensor_attrs_nchw(
+    std::shared_ptr<fe::graph::Tensor_attributes>& tensor,
+    int64_t uid,
+    const array& x) {
+  auto [shape, strides] = nhwc_to_nchw(x);
+  set_tensor_attrs(tensor, uid, x, shape, strides);
+}
+
+} // namespace mlx::core