[CUDA] Grouped mm

mlx/backend/cpu/masked_mm.cpp

@@ -549,4 +549,8 @@ void SegmentedMM::eval_cpu(const std::vector<array>& inputs, array& out) {
   });
 }
 
+void GroupedMM::eval_cpu(const std::vector<array>&, array&) {
+  throw std::runtime_error("[GroupedMM::eval_cpu] CPU grouped_mm NYI.");
+}
+
 } // namespace mlx::core

mlx/backend/cuda/CMakeLists.txt

@@ -36,6 +36,8 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/gemms/cublas_gemm.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/gemms/gather_gemm.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/gemms/grouped_gemm_unaligned.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/gemms/grouped_gemm.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/gemms/grouped_gemm_cudnn.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/jit_module.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp

mlx/backend/cuda/gemms/grouped_gemm.cpp

@@ -0,0 +1,24 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/gemms/grouped_gemm.h"
+#include "mlx/backend/cuda/cudnn_utils.h"
+
+#include <stdexcept>
+
+namespace mlx::core {
+
+void grouped_mm(
+    const array& a,
+    const array& b,
+    const array& offsets,
+    array& out,
+    cu::CommandEncoder& encoder) {
+#if CUDNN_VERSION >= 91800
+  cudnn_grouped_mm(a, b, offsets, out, encoder);
+#else
+  throw std::runtime_error(
+      "[grouped_mm] Grouped matmul requires cuDNN 9.18 or newer.");
+#endif
+}
+
+} // namespace mlx::core

mlx/backend/cuda/gemms/grouped_gemm.h

@@ -1,7 +1,6 @@
 // Copyright © 2025 Apple Inc.
 
 #pragma once
-
 namespace mlx::core {
 
 namespace cu {
@@ -36,4 +35,18 @@ void cutlass_segmented_mm(
     array& out,
     cu::CommandEncoder& encoder);
 
+void cudnn_grouped_mm(
+    const array& x,
+    const array& w,
+    const array& token_offsets,
+    array& out,
+    cu::CommandEncoder& encoder);
+
+void grouped_mm(
+    const array& a,
+    const array& b,
+    const array& offsets,
+    array& out,
+    cu::CommandEncoder& encoder);
+
 } // namespace mlx::core

mlx/backend/cuda/gemms/grouped_gemm_cudnn.cpp

@@ -0,0 +1,141 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/cudnn_utils.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/gemms/grouped_gemm.h"
+#include "mlx/backend/cuda/lru_cache.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/fast_primitives.h"
+
+#include <nvtx3/nvtx3.hpp>
+
+#include <optional>
+
+namespace mlx::core {
+
+#if CUDNN_VERSION >= 91800
+
+namespace {
+
+constexpr int GMM_NDIM = 3;
+
+struct GatherMMCacheKey {
+  int device_id;
+  fe::DataType_t cudnn_dtype;
+  int mode; // NONE / GATHER / SCATTER
+  std::array<int, GMM_NDIM> x_shape;
+  std::array<int64_t, GMM_NDIM> x_strides;
+  std::array<int, GMM_NDIM> w_shape;
+  std::array<int64_t, GMM_NDIM> w_strides;
+  std::array<int, GMM_NDIM> out_shape;
+};
+
+inline BytesKey<GatherMMCacheKey> build_grouped_mm_key(
+    cu::CommandEncoder& encoder,
+    const array& x,
+    const array& w,
+    const array& out) {
+  BytesKey<GatherMMCacheKey> key;
+  key.pod.device_id = encoder.device().cuda_device();
+  key.pod.cudnn_dtype = dtype_to_cudnn_type(x.dtype());
+  key.pod.mode = 0; // NONE
+  key.pod.x_shape = vector_key<GMM_NDIM>(x.shape());
+  key.pod.x_strides = vector_key<GMM_NDIM>(x.strides());
+  key.pod.w_shape = vector_key<GMM_NDIM>(w.shape());
+  key.pod.w_strides = vector_key<GMM_NDIM>(w.strides());
+  key.pod.out_shape = vector_key<GMM_NDIM>(out.shape());
+  return key;
+}
+
+enum UIDS { X, W, TOKEN_OFFSETS, TOKEN_INDEX, O };
+
+// cudnn expects specific shape and strides for grouped matmul:
+// [1, T, H]
+void set_moe_layout(
+    std::shared_ptr<fe::graph::Tensor_attributes>& t,
+    const array& x) {
+  int64_t L = x.shape(0);
+  int64_t D = x.shape(-1);
+  int64_t sL = x.strides(0);
+  int64_t sD = x.strides(-1);
+  t->set_dim({1, L, D}).set_stride({L * sL, sL, sD});
+}
+
+DnnGraph grouped_mm_graph(
+    cudnnHandle_t handle,
+    const array& x,
+    const array& w,
+    const array& token_offsets,
+    const array& output) {
+  DnnGraph graph(handle, x.dtype());
+
+  auto x_ = graph.tensor("X", X, x);
+  set_moe_layout(x_, x);
+  auto w_ = graph.tensor("W", W, w);
+  auto token_offsets_ =
+      graph.tensor("TOKEN_OFFSETS", TOKEN_OFFSETS, token_offsets);
+
+  auto moe_grouped_matmul_attr =
+      fe::graph::Moe_grouped_matmul_attributes()
+          .set_name("grouped_matmul")
+          .set_mode(fe::MoeGroupedMatmulMode_t::NONE);
+
+  std::shared_ptr<fe::graph::Tensor_attributes> token_index = nullptr;
+  std::shared_ptr<fe::graph::Tensor_attributes> token_ks = nullptr;
+
+  auto out_ = graph.moe_grouped_matmul(
+      x_, w_, token_offsets_, token_index, token_ks, moe_grouped_matmul_attr);
+  graph.tensor(out_, O, output);
+  set_moe_layout(out_, output);
+  out_->set_output(true);
+
+  CHECK_CUDNN_ERROR(graph.prepare());
+  graph.select_behavior_notes(
+      {fe::BehaviorNote_t::SUPPORTS_CUDA_GRAPH_NATIVE_API});
+  CHECK_CUDNN_ERROR(graph.build());
+  return graph;
+}
+
+auto& grouped_mm_cache() {
+  static thread_local LRUBytesKeyCache<GatherMMCacheKey, DnnGraph> cache(
+      "MLX_CUDA_GMM_CACHE_SIZE", /* default_capacity */ 256);
+  return cache;
+}
+
+} // namespace
+
+void cudnn_grouped_mm(
+    const array& x,
+    const array& w,
+    const array& token_offsets, // precomputed offsets for each expert
+    array& out,
+    cu::CommandEncoder& encoder) {
+  nvtx3::scoped_range r("cudnn_grouped_mm");
+
+  auto handle = get_cudnn_handle(encoder.device());
+
+  encoder.set_input_array(x);
+  encoder.set_input_array(w);
+  encoder.set_input_array(token_offsets);
+  encoder.set_output_array(out);
+
+  auto cache_key = build_grouped_mm_key(encoder, x, w, out);
+  auto& cache = grouped_mm_cache();
+  auto it = cache.find(cache_key);
+  if (it == cache.end()) {
+    auto graph = grouped_mm_graph(handle, x, w, token_offsets, out);
+    it = cache.emplace(cache_key, std::move(graph)).first;
+  }
+  auto& graph = it->second;
+
+  std::unordered_map<int64_t, void*> variant_pack{
+      {X, gpu_ptr<void>(x)},
+      {W, gpu_ptr<void>(w)},
+      {TOKEN_OFFSETS, gpu_ptr<void>(token_offsets)},
+      {O, gpu_ptr<void>(out)}};
+  CHECK_CUDNN_ERROR(graph.encode_graph(encoder, std::move(variant_pack)));
+}
+
+#endif
+
+} // namespace mlx::core

mlx/backend/cuda/matmul.cpp

@@ -39,7 +39,7 @@ ensure_batch_contiguous(const array& x, cu::CommandEncoder& encoder, Stream s) {
   }
 
   bool rc = true;
-  for (int i = 0; i < x.ndim() - 3; i++) {
+  for (int i = 0; i < static_cast<int>(x.ndim()) - 3; i++) {
     rc &= (x.strides(i + 1) * x.shape(i)) == x.strides(i);
   }
   if (rc) {
@@ -468,4 +468,30 @@ void SegmentedMM::eval_gpu(const std::vector<array>& inputs, array& out) {
       encoder);
 }
 
+void GroupedMM::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("GroupedMM::eval_gpu");
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  assert(inputs.size() == 3);
+  auto& a_pre = inputs[0];
+  auto& b_pre = inputs[1];
+  auto& offsets_pre = inputs[2];
+
+  if (out.size() == 0 || a_pre.size() == 0 || b_pre.size() == 0) {
+    array zero(0, a_pre.dtype());
+    encoder.add_temporary(zero);
+    fill_gpu(zero, out, s);
+    return;
+  }
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+
+  // a must be row contiguous
+  auto a = ensure_row_contiguous(a_pre, encoder, s);
+  auto b = std::get<2>(ensure_batch_contiguous(b_pre, encoder, s));
+  auto offsets = ensure_row_contiguous(offsets_pre, encoder, s);
+
+  grouped_mm(a, b, offsets, out, encoder);
+}
+
 } // namespace mlx::core

mlx/backend/metal/matmul.cpp

@@ -3098,4 +3098,8 @@ void SegmentedMM::eval_gpu(const std::vector<array>& inputs, array& out) {
   segmented_mm(a, b, segments, out, M, N, K, d, s);
 }
 
+void GroupedMM::eval_gpu(const std::vector<array>&, array&) {
+  throw std::runtime_error("[GroupedMM::eval_gpu] Metal grouped_mm NYI.");
+}
+
 } // namespace mlx::core

mlx/backend/no_cpu/primitives.cpp

@@ -74,6 +74,7 @@ NO_CPU(GatherQMM)
 NO_CPU(GatherQQMM)
 NO_CPU(Greater)
 NO_CPU(GreaterEqual)
+NO_CPU(GroupedMM)
 NO_CPU(Hadamard)
 NO_CPU(Imag)
 NO_CPU(Less)

mlx/backend/no_gpu/primitives.cpp

@@ -107,6 +107,7 @@ NO_GPU(GatherQMM)
 NO_GPU(GatherQQMM)
 NO_GPU(Greater)
 NO_GPU(GreaterEqual)
+NO_GPU(GroupedMM)
 NO_GPU(Hadamard)
 NO_GPU(Imag)
 NO_GPU(Less)

mlx/ops.cpp

@@ -6235,6 +6235,62 @@ array segmented_mm(
       {std::move(a), std::move(b), std::move(segments)});
 }
 
+array grouped_mm(
+    array a,
+    array b,
+    array token_offsets,
+    StreamOrDevice s /* = {} */) {
+  if (b.ndim() != 3) {
+    std::ostringstream msg;
+    msg << "[grouped_mm] Second input must have 3 dimensions (num_groups, N, K) but "
+        << "b.ndim() == " << b.ndim() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (a.ndim() != 2) {
+    std::ostringstream msg;
+    msg << "[grouped_mm] First input must have 2 dimensions but "
+        << "a.ndim() == " << a.ndim() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  if (a.shape(-1) != b.shape(-2)) {
+    std::ostringstream msg;
+    msg << "[grouped_mm] Last dimension of first input with shape " << a.shape()
+        << " must match second to last dimension of"
+        << " second input with shape " << b.shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto out_type = result_type(a, b);
+  if (!issubdtype(out_type, floating)) {
+    std::ostringstream msg;
+    msg << "[grouped_mm] Only real floating point types are supported but "
+        << a.dtype() << " and " << b.dtype()
+        << " were provided which results in " << out_type
+        << ", which is not a real floating point type.";
+    throw std::invalid_argument(msg.str());
+  }
+
+  a = astype(a, out_type, s);
+  b = astype(b, out_type, s);
+
+  if (!issubdtype(token_offsets.dtype(), integer)) {
+    throw std::invalid_argument(
+        "[grouped_mm] Got token_offsets with invalid dtype. Indices must be integral.");
+  }
+  token_offsets = astype(token_offsets, int32, s);
+
+  auto out_shape = a.shape();
+  out_shape.back() = b.shape(-1);
+
+  return array(
+      std::move(out_shape),
+      out_type,
+      std::make_shared<GroupedMM>(to_stream(s)),
+      {std::move(a), std::move(b), std::move(token_offsets)});
+}
+
 array diagonal(
     const array& a,
     int offset /* = 0 */,

mlx/ops.h

@@ -1702,6 +1702,14 @@ MLX_API array gather_mm(
 MLX_API array
 segmented_mm(array a, array b, array segments, StreamOrDevice s = {});
 
+/**
+ * Compute a matrix product over groups of rows. The rows of the first input
+ * are sorted by group and the offsets into each group are provided so that
+ * each group is multiplied by its corresponding matrix in the second input.
+ */
+MLX_API array
+grouped_mm(array a, array b, array token_offsets, StreamOrDevice s = {});
+
 /** Extract a diagonal or construct a diagonal array */
 MLX_API array diagonal(
     const array& a,

mlx/primitives.h

@@ -553,6 +553,16 @@ class GatherMM : public UnaryPrimitive {
   bool right_sorted_;
 };
 
+class GroupedMM : public UnaryPrimitive {
+ public:
+  explicit GroupedMM(Stream stream) : UnaryPrimitive(stream) {}
+
+  void eval_cpu(const std::vector<array>& inputs, array& out) override;
+  void eval_gpu(const std::vector<array>& inputs, array& out) override;
+
+  DEFINE_NAME(GroupedMM)
+};
+
 class SegmentedMM : public UnaryPrimitive {
  public:
   explicit SegmentedMM(Stream stream) : UnaryPrimitive(stream) {}