From 1d21d0e69687c2e0df322ce9d0031936e9230fef Mon Sep 17 00:00:00 2001 From: Cheng Date: Wed, 24 Dec 2025 09:42:56 +0900 Subject: [PATCH] [CUDA] Implement gather_mm_rhs (#2902) --- mlx/backend/cuda/CMakeLists.txt | 12 + mlx/backend/cuda/gemms/grouped_gemm.h | 25 ++ .../cuda/gemms/grouped_gemm_unaligned.cu | 288 ++++++++++++++++++ mlx/backend/cuda/matmul.cpp | 103 +++++++ mlx/backend/cuda/primitives.cpp | 1 - python/tests/cuda_skip.py | 2 +- python/tests/test_blas.py | 34 ++- 7 files changed, 458 insertions(+), 7 deletions(-) create mode 100644 mlx/backend/cuda/gemms/grouped_gemm.h create mode 100644 mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index 26ca4773..52fecee8 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/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 $) diff --git a/mlx/backend/cuda/gemms/grouped_gemm.h b/mlx/backend/cuda/gemms/grouped_gemm.h new file mode 100644 index 00000000..308a1ba9 --- /dev/null +++ b/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 diff --git a/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu b/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu new file mode 100644 index 00000000..0fe60dbd --- /dev/null +++ b/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 +#include +#include +#include +#include +#include + +namespace mlx::core { + +using ProblemSize = cutlass::gemm::GemmCoord; + +namespace cu { + +namespace cg = cooperative_groups; + +template +__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 +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; + 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; + + typename EpilogueOutputOp::Params epilogue_op( + /* alpha */ 1, /* beta */ 0); + typename GemmGrouped::Arguments args( + problem_sizes, + group_count, + GemmGrouped::sufficient(), + epilogue_op, + reinterpret_cast(a_ptrs), + reinterpret_cast(b_ptrs), + reinterpret_cast(out_ptrs), + reinterpret_cast(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(gemm_args); + int64_t* a_lds = gpu_ptr(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(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, + num_blocks, + block_dims, + group_count * sizeof(uint32_t), // sizeof(cum_histo) + gpu_ptr(indices), + indices.size(), + group_count, + a.shape(-1), // K + b.shape(-1), // N, + lda, + ldb, + out.itemsize(), + gpu_ptr(a), + gpu_ptr(b), + gpu_ptr(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; + switch (a.dtype()) { + case float32: + break; + case float16: + fun = grouped_gemm_v2; + break; + case bfloat16: + fun = grouped_gemm_v2; + 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 diff --git a/mlx/backend/cuda/matmul.cpp b/mlx/backend/cuda/matmul.cpp index 392a13ad..66b3a22a 100644 --- a/mlx/backend/cuda/matmul.cpp +++ b/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 +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& inputs, array& out) { @@ -254,4 +321,40 @@ void AddMM::eval_gpu(const std::vector& inputs, array& out) { beta_); } +void GatherMM::eval_gpu(const std::vector& 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 diff --git a/mlx/backend/cuda/primitives.cpp b/mlx/backend/cuda/primitives.cpp index 6461c4ec..56e5969e 100644 --- a/mlx/backend/cuda/primitives.cpp +++ b/mlx/backend/cuda/primitives.cpp @@ -33,7 +33,6 @@ void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { NO_GPU(BlockMaskedMM) NO_GPU(FFT) -NO_GPU(GatherMM) NO_GPU(GatherQMM) NO_GPU(Hadamard) NO_GPU_MULTI(LUF) diff --git a/python/tests/cuda_skip.py b/python/tests/cuda_skip.py index f9e6fe8e..caa7213c 100644 --- a/python/tests/cuda_skip.py +++ b/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 diff --git a/python/tests/test_blas.py b/python/tests/test_blas.py index 8e97e1f4..469d2407 100644 --- a/python/tests/test_blas.py +++ b/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],