diff --git a/mlx/backend/cpu/masked_mm.cpp b/mlx/backend/cpu/masked_mm.cpp index db38f39424..a4983341d0 100644 --- a/mlx/backend/cpu/masked_mm.cpp +++ b/mlx/backend/cpu/masked_mm.cpp @@ -549,4 +549,8 @@ void SegmentedMM::eval_cpu(const std::vector& inputs, array& out) { }); } +void GroupedMM::eval_cpu(const std::vector&, array&) { + throw std::runtime_error("[GroupedMM::eval_cpu] CPU grouped_mm NYI."); +} + } // namespace mlx::core diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index 72b5868004..25e680905b 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/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 diff --git a/mlx/backend/cuda/gemms/grouped_gemm.cpp b/mlx/backend/cuda/gemms/grouped_gemm.cpp new file mode 100644 index 0000000000..c22dbff202 --- /dev/null +++ b/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 + +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 diff --git a/mlx/backend/cuda/gemms/grouped_gemm.h b/mlx/backend/cuda/gemms/grouped_gemm.h index 844b8f440e..bdf6db641f 100644 --- a/mlx/backend/cuda/gemms/grouped_gemm.h +++ b/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 diff --git a/mlx/backend/cuda/gemms/grouped_gemm_cudnn.cpp b/mlx/backend/cuda/gemms/grouped_gemm_cudnn.cpp new file mode 100644 index 0000000000..c1b79f5b2e --- /dev/null +++ b/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 + +#include + +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 x_shape; + std::array x_strides; + std::array w_shape; + std::array w_strides; + std::array out_shape; +}; + +inline BytesKey build_grouped_mm_key( + cu::CommandEncoder& encoder, + const array& x, + const array& w, + const array& out) { + BytesKey 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(x.shape()); + key.pod.x_strides = vector_key(x.strides()); + key.pod.w_shape = vector_key(w.shape()); + key.pod.w_strides = vector_key(w.strides()); + key.pod.out_shape = vector_key(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& 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 token_index = nullptr; + std::shared_ptr 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 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 variant_pack{ + {X, gpu_ptr(x)}, + {W, gpu_ptr(w)}, + {TOKEN_OFFSETS, gpu_ptr(token_offsets)}, + {O, gpu_ptr(out)}}; + CHECK_CUDNN_ERROR(graph.encode_graph(encoder, std::move(variant_pack))); +} + +#endif + +} // namespace mlx::core diff --git a/mlx/backend/cuda/matmul.cpp b/mlx/backend/cuda/matmul.cpp index 7505ea2cba..f48bdd2182 100644 --- a/mlx/backend/cuda/matmul.cpp +++ b/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(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& inputs, array& out) { encoder); } +void GroupedMM::eval_gpu(const std::vector& 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 diff --git a/mlx/backend/metal/matmul.cpp b/mlx/backend/metal/matmul.cpp index fc81e72cca..66671b46e7 100644 --- a/mlx/backend/metal/matmul.cpp +++ b/mlx/backend/metal/matmul.cpp @@ -3098,4 +3098,8 @@ void SegmentedMM::eval_gpu(const std::vector& inputs, array& out) { segmented_mm(a, b, segments, out, M, N, K, d, s); } +void GroupedMM::eval_gpu(const std::vector&, array&) { + throw std::runtime_error("[GroupedMM::eval_gpu] Metal grouped_mm NYI."); +} + } // namespace mlx::core diff --git a/mlx/backend/no_cpu/primitives.cpp b/mlx/backend/no_cpu/primitives.cpp index e522307156..605128dd70 100644 --- a/mlx/backend/no_cpu/primitives.cpp +++ b/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) diff --git a/mlx/backend/no_gpu/primitives.cpp b/mlx/backend/no_gpu/primitives.cpp index f17f12cfaf..a53b6faf55 100644 --- a/mlx/backend/no_gpu/primitives.cpp +++ b/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) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index be368e8720..e2b36beb31 100644 --- a/mlx/ops.cpp +++ b/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(to_stream(s)), + {std::move(a), std::move(b), std::move(token_offsets)}); +} + array diagonal( const array& a, int offset /* = 0 */, diff --git a/mlx/ops.h b/mlx/ops.h index f597753b1e..fc83dbc780 100644 --- a/mlx/ops.h +++ b/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, diff --git a/mlx/primitives.h b/mlx/primitives.h index 0cfc71bf04..ade3dd1b7a 100644 --- a/mlx/primitives.h +++ b/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& inputs, array& out) override; + void eval_gpu(const std::vector& inputs, array& out) override; + + DEFINE_NAME(GroupedMM) +}; + class SegmentedMM : public UnaryPrimitive { public: explicit SegmentedMM(Stream stream) : UnaryPrimitive(stream) {} diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 414109739a..131e218137 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -4944,6 +4944,28 @@ void init_ops(nb::module_& m) { Returns: array: The result per segment of shape ``MxN``. )pbdoc"); + m.def( + "grouped_mm", + &mx::grouped_mm, + nb::arg(), + nb::arg(), + "token_offsets"_a, + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def grouped_mm(a: array, b: array, /, *, token_offsets: array, stream: StreamOrDevice = None) -> array"), + R"pbdoc( + Perform a matrix multiplication but segment the inner dimension and + save the result for each segment separately. + + Args: + a (array): Input array of shape ``MxK``. + b (array): Input array of shape ``ExKxN``. + token_offsets (array): The offsets into the inner dimension for each segment. + + Returns: + array: The result of shape ``MxN``. + )pbdoc"); m.def( "tensordot", [](const mx::array& a, diff --git a/python/tests/test_blas.py b/python/tests/test_blas.py index 7dbbaea049..5c506a85c2 100644 --- a/python/tests/test_blas.py +++ b/python/tests/test_blas.py @@ -1,8 +1,9 @@ # Copyright © 2023-2024 Apple Inc. import math +import os import unittest -from itertools import permutations +from itertools import product import mlx.core as mx import mlx_tests @@ -1609,6 +1610,61 @@ def segmented_mm_ref(a, b, s): c = mx.segmented_mm(a, a.T, s) self.assertEqual(c.shape, (2, 2, 4, 10, 10)) + @unittest.skipIf("CI" in os.environ, "No device that supports grouped_mm in CI") + def test_grouped_mm(self): + def token_offsets(sizes): + offsets = [sum(sizes[:g]) for g in range(len(sizes))] + return mx.array(offsets, mx.int32).reshape(len(sizes), 1, 1) + + def grouped_mm_ref(a, b, sizes): + c = [] + lo = 0 + for g, size in enumerate(sizes): + if size > 0: + c.append(a[lo : lo + size] @ b[g]) + lo += size + return mx.concatenate(c, axis=0) + + K, N = 64, 32 + allocations = [ + [16, 16, 16, 16], + [5, 11, 1, 47], + [5, 11, 1, 47, 15], + [0, 7, 9, 0], + [0, 0, 0, 24], + [64], + ] + # in cudnn's grouped matmul tf32 can't be disabled. + # So we only test float16 and bfloat16 for now. + dtypes = [(mx.float16, 1e-3), (mx.bfloat16, 1e-2)] + + for allocation, b_transposed, a_transposed, (dtype, tol) in product( + allocations, (True, False), (True, False), dtypes + ): + with self.subTest( + sizes=allocation, + b_transposed=b_transposed, + a_transposed=a_transposed, + dtype=dtype, + ): + E = len(allocation) + T = sum(allocation) + a_shape = (K, T) if a_transposed else (T, K) + b_shape = (E, N, K) if b_transposed else (E, K, N) + a = mx.random.normal(a_shape, dtype=dtype) + b = mx.random.normal(b_shape, dtype=dtype) + if a_transposed: + a = a.swapaxes(-1, -2) + if b_transposed: + b = b.swapaxes(-1, -2) + offsets = token_offsets(allocation) + + c1 = grouped_mm_ref(a, b, allocation) + c2 = mx.grouped_mm(a, b, token_offsets=offsets) + self.assertEqual(c2.shape, (T, N)) + self.assertEqual(c2.dtype, dtype) + self.assertTrue(mx.allclose(c1, c2, rtol=tol, atol=tol)) + def test_gemv_gemm_same_precision(self): mx.random.seed(0) N = 256