[CUDA] Implement gather_mm_rhs (#2902)

mlx/backend/cuda/CMakeLists.txt

@@ -29,6 +29,7 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/fence.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/gemms/gemv.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/gemms/cublas_gemm.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/gemms/grouped_gemm_unaligned.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/jit_module.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/kernel_utils.cu
@@ -226,3 +227,14 @@ target_link_libraries(mlx PRIVATE cudnn_frontend)
 # Link with the actual cuDNN libraries.
 include(${cudnn_frontend_SOURCE_DIR}/cmake/cuDNN.cmake)
 target_link_libraries(mlx PRIVATE CUDNN::cudnn_all)
+
+# Use header-only CUTLASS.
+FetchContent_Declare(
+  cutlass
+  GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
+  GIT_TAG v4.3.2
+  GIT_SHALLOW TRUE
+  SOURCE_SUBDIR include EXCLUDE_FROM_ALL)
+FetchContent_MakeAvailable(cutlass)
+target_include_directories(
+  mlx PRIVATE $<BUILD_INTERFACE:${cutlass_SOURCE_DIR}/include>)

mlx/backend/cuda/gemms/grouped_gemm.h

@@ -0,0 +1,25 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+namespace mlx::core {
+
+namespace cu {
+class CommandEncoder;
+}
+
+class array;
+
+void cutlass_grouped_gemm_unaligned(
+    bool a_transposed,
+    int lda,
+    bool b_transposed,
+    int ldb,
+    int group_count,
+    const array& a,
+    const array& b,
+    const array& indices,
+    array& out,
+    cu::CommandEncoder& encoder);
+
+} // namespace mlx::core

mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu

@@ -0,0 +1,288 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/gemms/grouped_gemm.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/dtype_utils.h"
+
+#include <cooperative_groups.h>
+#include <cutlass/gemm/device/default_gemm_configuration.h>
+#include <cutlass/gemm/device/gemm_grouped.h>
+#include <cutlass/gemm/kernel/default_gemm_grouped.h>
+#include <fmt/format.h>
+#include <nvtx3/nvtx3.hpp>
+
+namespace mlx::core {
+
+using ProblemSize = cutlass::gemm::GemmCoord;
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <int N_READS>
+__global__ void prepare_grouped_mm_data(
+    const uint32_t* indices,
+    size_t size,
+    int group_count,
+    int K,
+    int N,
+    int lda,
+    int ldb,
+    int item_size,
+    int8_t* a_start,
+    int8_t* b_start,
+    int8_t* out_start,
+    int a_batch_stride,
+    int b_batch_stride,
+    int out_batch_stride,
+    ProblemSize* problem_sizes,
+    int64_t* a_lds,
+    int64_t* b_lds,
+    int64_t* out_lds,
+    void** a_ptrs,
+    void** b_ptrs,
+    void** out_ptrs) {
+  auto block = cg::this_thread_block();
+
+  // cumsum(histogram(indices)) - offset for each group.
+  extern __shared__ uint32_t cum_histo[];
+
+  int group = block.thread_rank();
+  if (group < group_count) {
+    cum_histo[group] = 0;
+  }
+
+  block.sync();
+
+  // Since |indices| is sorted, the position where element changes would be its
+  // cumulative histogram.
+  size_t elems_per_block = block.num_threads() * N_READS;
+  for (int r = 0; r < cuda::ceil_div(size, elems_per_block); ++r) {
+    // TODO: Use vectorized read.
+    for (int i = 0; i < N_READS; ++i) {
+      size_t pos = r * elems_per_block + group * N_READS + i;
+      if (pos >= size) {
+        break;
+      }
+      auto elem = indices[pos];
+      auto next = pos < size - 1 ? indices[pos + 1] : group_count;
+      while (elem < next) {
+        cum_histo[elem] = pos + 1;
+        elem++;
+      }
+    }
+  }
+
+  block.sync();
+
+  if (group < group_count) {
+    // Fill shapes.
+    int delta =
+        group == 0 ? cum_histo[0] : cum_histo[group] - cum_histo[group - 1];
+    problem_sizes[group] = {delta, N, K};
+    a_lds[group] = lda;
+    b_lds[group] = ldb;
+    out_lds[group] = N;
+    // Fill pointers.
+    auto offset = group == 0 ? 0 : cum_histo[group - 1];
+    a_ptrs[group] = a_start + offset * item_size * a_batch_stride;
+    b_ptrs[group] = b_start + group * item_size * b_batch_stride;
+    out_ptrs[group] = out_start + offset * item_size * out_batch_stride;
+  }
+}
+
+} // namespace cu
+
+namespace {
+
+template <typename T, int kAlignment, typename Arch, typename OpClass>
+void grouped_gemm_v2(
+    bool a_transposed,
+    bool b_transposed,
+    int group_count,
+    ProblemSize* problem_sizes,
+    int64_t* a_lds,
+    int64_t* b_lds,
+    int64_t* out_lds,
+    void* a_ptrs,
+    void* b_ptrs,
+    void* out_ptrs,
+    cu::CommandEncoder& encoder) {
+  using ElementAccumulator = float;
+  using GemmConfiguration = typename cutlass::gemm::device::
+      DefaultGemmConfiguration<OpClass, Arch, T, T, T, ElementAccumulator>;
+  using EpilogueOutputOp = typename GemmConfiguration::EpilogueOutputOp;
+
+  dispatch_bool(a_transposed, [&](auto a_transposed_tag) {
+    dispatch_bool(b_transposed, [&](auto b_transposed_tag) {
+      using LayoutA = std::conditional_t<
+          a_transposed_tag,
+          cutlass::layout::ColumnMajor,
+          cutlass::layout::RowMajor>;
+      using LayoutB = std::conditional_t<
+          b_transposed_tag,
+          cutlass::layout::ColumnMajor,
+          cutlass::layout::RowMajor>;
+      using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmGrouped<
+          T,
+          LayoutA,
+          cutlass::ComplexTransform::kNone,
+          kAlignment,
+          T,
+          LayoutB,
+          cutlass::ComplexTransform::kNone,
+          kAlignment,
+          T,
+          cutlass::layout::RowMajor,
+          ElementAccumulator,
+          OpClass,
+          Arch,
+          typename GemmConfiguration::ThreadblockShape,
+          typename GemmConfiguration::WarpShape,
+          typename GemmConfiguration::InstructionShape,
+          EpilogueOutputOp,
+          cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,
+          GemmConfiguration::kStages>::GemmKernel;
+      using GemmGrouped =
+          typename cutlass::gemm::device::GemmGrouped<GemmKernel>;
+
+      typename EpilogueOutputOp::Params epilogue_op(
+          /* alpha */ 1, /* beta */ 0);
+      typename GemmGrouped::Arguments args(
+          problem_sizes,
+          group_count,
+          GemmGrouped::sufficient(),
+          epilogue_op,
+          reinterpret_cast<T**>(a_ptrs),
+          reinterpret_cast<T**>(b_ptrs),
+          reinterpret_cast<T**>(out_ptrs),
+          reinterpret_cast<T**>(out_ptrs),
+          a_lds,
+          b_lds,
+          out_lds,
+          out_lds);
+
+      GemmGrouped gemm;
+      cutlass::Status status = gemm.initialize(args, nullptr, encoder.stream());
+      if (status != cutlass::Status::kSuccess) {
+        throw std::runtime_error(fmt::format(
+            "Failed to initialize GemmGrouped: {}",
+            cutlass::cutlassGetStatusString(status)));
+      }
+
+      auto capture = encoder.capture_context();
+      status = gemm.run(encoder.stream());
+      if (status != cutlass::Status::kSuccess) {
+        throw std::runtime_error(fmt::format(
+            "Failed to run GemmGrouped: {}",
+            cutlass::cutlassGetStatusString(status)));
+      }
+    });
+  });
+}
+
+} // namespace
+
+void cutlass_grouped_gemm_unaligned(
+    bool a_transposed,
+    int lda,
+    bool b_transposed,
+    int ldb,
+    int group_count,
+    const array& a,
+    const array& b,
+    const array& indices,
+    array& out,
+    cu::CommandEncoder& encoder) {
+  // Prepare device pointers for matmul.
+  int problem_sizes_nbytes =
+      group_count * cuda::ceil_div(sizeof(ProblemSize), 8) * 8;
+  int nbytes = problem_sizes_nbytes +
+      group_count * (3 * sizeof(void*) + 3 * sizeof(int64_t));
+  nbytes = cuda::ceil_div(nbytes, 256) * 256;
+  array gemm_args(cu::malloc_async(nbytes, encoder), {nbytes}, int8);
+  encoder.add_temporary(gemm_args);
+
+  ProblemSize* problem_sizes = gpu_ptr<ProblemSize>(gemm_args);
+  int64_t* a_lds = gpu_ptr<int64_t>(gemm_args) + problem_sizes_nbytes / 8;
+  int64_t* b_lds = a_lds + group_count;
+  int64_t* out_lds = b_lds + group_count;
+  void** a_ptrs = reinterpret_cast<void**>(out_lds + group_count);
+  void** b_ptrs = a_ptrs + group_count;
+  void** out_ptrs = b_ptrs + group_count;
+
+  // Fill the pointers by computing offsets from indices.
+  constexpr int N_READS = 4;
+  size_t n_threads = cuda::ceil_div(indices.size(), N_READS);
+  n_threads = group_count < n_threads ? n_threads : group_count;
+  dim3 block_dims(std::min(n_threads, 1024ul));
+  dim3 num_blocks(1);
+
+  encoder.set_input_array(indices);
+  encoder.set_output_array(gemm_args);
+  encoder.add_kernel_node(
+      cu::prepare_grouped_mm_data<N_READS>,
+      num_blocks,
+      block_dims,
+      group_count * sizeof(uint32_t), // sizeof(cum_histo)
+      gpu_ptr<uint32_t>(indices),
+      indices.size(),
+      group_count,
+      a.shape(-1), // K
+      b.shape(-1), // N,
+      lda,
+      ldb,
+      out.itemsize(),
+      gpu_ptr<int8_t>(a),
+      gpu_ptr<int8_t>(b),
+      gpu_ptr<int8_t>(out),
+      a.shape(-2) * a.shape(-1), // a_batch_stride
+      b.shape(-2) * b.shape(-1), // b_batch_stride
+      out.shape(-2) * out.shape(-1), // out_batch_stride
+      problem_sizes,
+      a_lds,
+      b_lds,
+      out_lds,
+      a_ptrs,
+      b_ptrs,
+      out_ptrs);
+
+  // Invoke grouped GEMM.
+  constexpr int kAlignment = 1;
+  using Arch = cutlass::arch::Sm75;
+  using OpClass = cutlass::arch::OpClassSimt;
+  auto* fun = grouped_gemm_v2<float, kAlignment, Arch, OpClass>;
+  switch (a.dtype()) {
+    case float32:
+      break;
+    case float16:
+      fun = grouped_gemm_v2<cutlass::half_t, kAlignment, Arch, OpClass>;
+      break;
+    case bfloat16:
+      fun = grouped_gemm_v2<cutlass::bfloat16_t, kAlignment, Arch, OpClass>;
+      break;
+    default:
+      throw std::runtime_error(fmt::format(
+          "Unsupported dtype in cutlass_grouped_gemm_sm75: {}.",
+          dtype_to_string(a.dtype())));
+  }
+
+  encoder.set_input_array(a);
+  encoder.set_input_array(b);
+  encoder.set_input_array(gemm_args);
+  encoder.set_output_array(out);
+  fun(a_transposed,
+      b_transposed,
+      group_count,
+      problem_sizes,
+      a_lds,
+      b_lds,
+      out_lds,
+      a_ptrs,
+      b_ptrs,
+      out_ptrs,
+      encoder);
+}
+
+} // namespace mlx::core

mlx/backend/cuda/matmul.cpp

@@ -4,6 +4,7 @@
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/gemms/cublas_gemm.h"
 #include "mlx/backend/cuda/gemms/gemv.h"
+#include "mlx/backend/cuda/gemms/grouped_gemm.h"
 #include "mlx/backend/gpu/copy.h"
 #include "mlx/primitives.h"
 
@@ -29,6 +30,38 @@ check_transpose(cu::CommandEncoder& enc, const Stream& s, const array& arr) {
   }
 }
 
+std::tuple<bool, int64_t, array>
+ensure_batch_contiguous(const array& x, cu::CommandEncoder& encoder, Stream s) {
+  if (x.flags().row_contiguous) {
+    return std::make_tuple(false, x.strides(-2), x);
+  }
+
+  bool rc = true;
+  for (int i = 0; i < x.ndim() - 3; i++) {
+    rc &= (x.strides(i + 1) * x.shape(i)) == x.strides(i);
+  }
+  if (rc) {
+    return check_transpose(encoder, s, x);
+  }
+
+  array x_copy = contiguous_copy_gpu(x, s);
+  encoder.add_temporary(x_copy);
+  return std::make_tuple(false, x_copy.strides(-2), x_copy);
+}
+
+array ensure_row_contiguous(
+    const array& x,
+    cu::CommandEncoder& encoder,
+    Stream s) {
+  if (!x.flags().row_contiguous) {
+    array x_copy = contiguous_copy_gpu(x, s);
+    encoder.add_temporary(x_copy);
+    return x_copy;
+  } else {
+    return x;
+  }
+}
+
 void gemm_and_bias(
     cu::CommandEncoder& encoder,
     int M,
@@ -103,6 +136,40 @@ void gemm_and_bias(
       encoder, out, a, b, batch_shape, a_batch_strides, b_batch_strides, alpha);
 }
 
+void gather_mm_rhs(
+    const array& a_,
+    const array& b_,
+    const array& indices_,
+    array& out,
+    cu::CommandEncoder& encoder,
+    Stream s) {
+  if (a_.size() / a_.shape(-2) / a_.shape(-1) != indices_.size()) {
+    throw std::runtime_error("[gather_mm] Broadcasting lhs is not supported.");
+  }
+
+  int group_count = b_.size() / b_.shape(-1) / b_.shape(-2);
+  if (group_count > 1024) {
+    throw std::runtime_error(
+        "[gather_mm] Group count can not be larger than 1024.");
+  }
+
+  auto [a_transposed, lda, a] = ensure_batch_contiguous(a_, encoder, s);
+  auto [b_transposed, ldb, b] = ensure_batch_contiguous(b_, encoder, s);
+  auto indices = ensure_row_contiguous(indices_, encoder, s);
+
+  cutlass_grouped_gemm_unaligned(
+      a_transposed,
+      lda,
+      b_transposed,
+      ldb,
+      group_count,
+      a,
+      b,
+      indices,
+      out,
+      encoder);
+}
+
 } // namespace
 
 void Matmul::eval_gpu(const std::vector<array>& inputs, array& out) {
@@ -254,4 +321,40 @@ void AddMM::eval_gpu(const std::vector<array>& inputs, array& out) {
       beta_);
 }
 
+void GatherMM::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("GatherMM::eval_gpu");
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  assert(inputs.size() == 4);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  auto& lhs_indices = inputs[2];
+  auto& rhs_indices = inputs[3];
+
+  // Return 0s if either input is empty.
+  if (a.size() == 0 || b.size() == 0) {
+    array zero(0, a.dtype());
+    encoder.add_temporary(zero);
+    fill_gpu(zero, out, s);
+    return;
+  }
+
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+
+  // Extract shapes from inputs.
+  int M = a.shape(-2);
+  int N = b.shape(-1);
+  int K = a.shape(-1);
+
+  // We are walking a in order and b is also in order so we can batch up the
+  // matmuls and reuse reading a and b.
+  if (M == 1 && right_sorted_ == true) {
+    gather_mm_rhs(a, b, rhs_indices, out, encoder, s);
+    return;
+  }
+
+  throw std::runtime_error("NYI");
+}
+
 } // namespace mlx::core

mlx/backend/cuda/primitives.cpp

@@ -33,7 +33,6 @@ void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
 
 NO_GPU(BlockMaskedMM)
 NO_GPU(FFT)
-NO_GPU(GatherMM)
 NO_GPU(GatherQMM)
 NO_GPU(Hadamard)
 NO_GPU_MULTI(LUF)

python/tests/cuda_skip.py

@@ -6,7 +6,7 @@ cuda_skip = {
     # Gather matmul NYI
     "TestBlas.test_gather_matmul",
     "TestBlas.test_gather_matmul_grad",
-    "TestBlas.test_gather_mm_sorted",
+    "TestBlas.test_gather_mm_sorted_vjp",
     # Segmented matmul NYI
     "TestBlas.test_segmented_mm",
     # Hadamard NYI

python/tests/test_blas.py

@@ -1235,15 +1235,39 @@ class TestBlas(mlx_tests.MLXTestCase):
         def gather_mm_test(a, b, rhs):
             return mx.gather_mm(a, b, rhs_indices=rhs, sorted_indices=True)
 
+        dtypes = [(mx.float32, 1e-4)]
+        if mx.cuda.is_available():
+            dtypes += [
+                (mx.float16, 1e-3),
+                (mx.bfloat16, 1e-2),
+            ]
+
+        for b_transposed in (True, False):
+            for dtype, tol in dtypes:
+                with self.subTest(b_transposed=b_transposed, dtype=dtype):
+                    a = mx.random.normal((100, 1, 100), dtype=dtype)
+                    b = mx.random.normal((8, 100, 100), dtype=dtype)
+                    if b_transposed:
+                        b = b.swapaxes(-1, -2)
+                    rhs = mx.sort(mx.random.randint(0, 8, shape=(100,)))
+
+                    c1 = gather_mm_ref(a, b, rhs)
+                    c2 = gather_mm_test(a, b, rhs)
+                    self.assertTrue(mx.allclose(c1, c2, rtol=tol, atol=tol))
+
+    def test_gather_mm_sorted_vjp(self):
+        def gather_mm_ref(a, b, rhs):
+            b = b[rhs]
+            return a @ b
+
+        def gather_mm_test(a, b, rhs):
+            return mx.gather_mm(a, b, rhs_indices=rhs, sorted_indices=True)
+
         a = mx.random.normal((100, 1, 100))
         b = mx.random.normal((8, 100, 100))
         rhs = mx.sort(mx.random.randint(0, 8, shape=(100,)))
 
-        c1 = gather_mm_ref(a, b, rhs)
-        c2 = gather_mm_test(a, b, rhs)
-        self.assertTrue(mx.allclose(c1, c2, atol=1e-4))
-
-        cotan = mx.random.normal(c1.shape)
+        cotan = mx.random.normal((100, 1, 100))
         c1, dc1 = mx.vjp(
             lambda a, b: gather_mm_ref(a, b, rhs),
             [a, b],