From becc7690128df1fe12e4e62b4581569a7ce187f2 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Thu, 22 Jan 2026 17:20:48 -0800 Subject: [PATCH] CUDA gather mv (#3039) --- mlx/backend/cuda/gemms/gemv.cu | 123 +++++++++++++++++++++++++++++++++ mlx/backend/cuda/gemms/gemv.h | 10 +++ mlx/backend/cuda/matmul.cpp | 13 ++++ 3 files changed, 146 insertions(+) diff --git a/mlx/backend/cuda/gemms/gemv.cu b/mlx/backend/cuda/gemms/gemv.cu index 6966a583..daa27522 100644 --- a/mlx/backend/cuda/gemms/gemv.cu +++ b/mlx/backend/cuda/gemms/gemv.cu @@ -104,6 +104,68 @@ __global__ void gemv_batched( mat + mat_offset, vec + vec_offset, out + batch_idx * rows, rows, cols); } +template +__global__ void gemv_gather( + const T* mat, + const T* vec, + T* out, + uint32_t* mat_indices, + uint32_t* vec_indices, + int rows, + int cols, + const __grid_constant__ Shape mat_batch_shape, + const __grid_constant__ Strides mat_batch_strides, + int mat_batch_ndim, + const __grid_constant__ Shape vec_batch_shape, + const __grid_constant__ Strides vec_batch_strides, + int vec_batch_ndim, + const __grid_constant__ Shape index_shape, + const __grid_constant__ Strides mat_index_strides, + const __grid_constant__ Strides vec_index_strides, + int index_batch_ndim) { + auto block = cg::this_thread_block(); + auto indices_idx = block.group_index().y; + uint32_t index_mat, index_vec; + if (index_batch_ndim > 1) { + auto [mat_idx_offset, vec_idx_offset] = elem_to_loc( + indices_idx, + index_shape.data(), + mat_index_strides.data(), + vec_index_strides.data(), + index_batch_ndim); + index_mat = mat_indices[mat_idx_offset]; + index_vec = vec_indices[vec_idx_offset]; + } else { + index_mat = mat_indices[indices_idx * mat_index_strides[0]]; + index_vec = vec_indices[indices_idx * vec_index_strides[0]]; + } + + int64_t mat_offset; + if (mat_batch_ndim > 1) { + mat_offset = elem_to_loc( + index_mat, + mat_batch_shape.data(), + mat_batch_strides.data(), + mat_batch_ndim); + } else { + mat_offset = index_mat * mat_batch_strides[0]; + } + + int64_t vec_offset; + if (vec_batch_ndim > 1) { + vec_offset = elem_to_loc( + index_vec, + vec_batch_shape.data(), + vec_batch_strides.data(), + vec_batch_ndim); + } else { + vec_offset = index_vec * vec_batch_strides[0]; + } + + gemv_impl( + mat + mat_offset, vec + vec_offset, out + indices_idx * rows, rows, cols); +} + bool can_use_gemv(int M, int N, int K, bool a_transposed, bool b_transposed) { return K % 32 == 0 && ((M == 1 && b_transposed) || (N == 1 && !a_transposed)); } @@ -201,4 +263,65 @@ void gemv( }); } +void gather_mv( + const array& mat_, + const array& vec_, + const array& mat_indices, + const array& vec_indices, + array& out, + int N, + int K, + CommandEncoder& encoder) { + encoder.set_input_array(mat_); + encoder.set_input_array(vec_); + encoder.set_input_array(mat_indices); + encoder.set_input_array(vec_indices); + encoder.set_output_array(out); + dispatch_inexact_types(out.dtype(), "gather_mv", [&](auto type_tag) { + using DataType = cuda_type_t; + dim3 block_dims{WARP_SIZE, rows_per_block}; + int rows = N; + int cols = K; + uint32_t batch_size = static_cast(out.size() / N); + const DataType* mat = gpu_ptr(mat_); + const DataType* vec = gpu_ptr(vec_); + + uint32_t num_blocks_x = (rows + rows_per_block - 1) / rows_per_block; + int n_per_t; + if (K % 128 == 0 && is_aligned<4>(mat) && is_aligned<4>(vec)) { + n_per_t = 4; + } else if (K % 64 == 0 && is_aligned<2>(mat) && is_aligned<2>(vec)) { + n_per_t = 2; + } else { + n_per_t = 1; + } + + dispatch_n_per_thread(n_per_t, [&](auto n_per_thread) { + auto kernel = gemv_gather; + encoder.add_kernel_node( + kernel, + dim3{num_blocks_x, batch_size}, + block_dims, + 0, + mat, + vec, + gpu_ptr(out), + gpu_ptr(mat_indices), + gpu_ptr(vec_indices), + rows, + cols, + const_param(mat_.shape()), + const_param(mat_.strides()), + mat_.ndim() - 2, + const_param(vec_.shape()), + const_param(vec_.strides()), + vec_.ndim() - 2, + const_param(mat_indices.shape()), + const_param(mat_indices.strides()), + const_param(vec_indices.strides()), + mat_indices.ndim()); + }); + }); +} + } // namespace mlx::core::cu diff --git a/mlx/backend/cuda/gemms/gemv.h b/mlx/backend/cuda/gemms/gemv.h index 27173aad..b0dd49d4 100644 --- a/mlx/backend/cuda/gemms/gemv.h +++ b/mlx/backend/cuda/gemms/gemv.h @@ -21,4 +21,14 @@ void gemv( const mlx::core::Strides& b_batch_strides, CommandEncoder& encoder); +void gather_mv( + const array& mat, + const array& vec, + const array& mat_indices, + const array& vec_indices, + array& out, + int N, + int K, + CommandEncoder& encoder); + } // namespace mlx::core::cu diff --git a/mlx/backend/cuda/matmul.cpp b/mlx/backend/cuda/matmul.cpp index 66b3a22a..725f1b40 100644 --- a/mlx/backend/cuda/matmul.cpp +++ b/mlx/backend/cuda/matmul.cpp @@ -354,6 +354,19 @@ void GatherMM::eval_gpu(const std::vector& inputs, array& out) { return; } + auto [transposed_a, lda, a_] = check_transpose(encoder, s, a); + auto [transposed_b, ldb, b_] = check_transpose(encoder, s, b); + auto use_gemv = cu::can_use_gemv(M, N, K, transposed_a, transposed_b); + if (M == 1 && use_gemv) { + gather_mv(b_, a_, rhs_indices, lhs_indices, out, N, K, encoder); + return; + } + + if (N == 1 && use_gemv) { + gather_mv(a_, b_, lhs_indices, rhs_indices, out, M, K, encoder); + return; + } + throw std::runtime_error("NYI"); }