From 6189e6fc2ce847198da85979f8caddf3a677dc89 Mon Sep 17 00:00:00 2001 From: Mohamed Elashri Date: Tue, 15 Sep 2026 12:39:29 +0200 Subject: [PATCH] cuda: support row-contiguous SUM_ROWS (llama/26308) * cuda: support row-contiguous SUM_ROWS * organize the code and add GGML_OP_MEAN to support row-contiguous tensors using the same shared kernel, and add a test to MEAN permute/slice * Keep original comments and add if/else branch --- ggml/src/ggml-cuda/ggml-cuda.cu | 2 ++ ggml/src/ggml-cuda/mean.cu | 21 ++++++++++----- ggml/src/ggml-cuda/reduce_rows.cuh | 43 +++++++++++++++++++++++++----- ggml/src/ggml-cuda/sumrows.cu | 20 +++++++++----- 4 files changed, 65 insertions(+), 21 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 43003245c..74bb47145 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5466,7 +5466,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g return true; #endif case GGML_OP_SUM_ROWS: + return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && ggml_is_contiguous_rows(op->src[0]); case GGML_OP_MEAN: + return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && ggml_is_contiguous_rows(op->src[0]); case GGML_OP_GROUP_NORM: return ggml_is_contiguous(op->src[0]); case GGML_OP_PAD: diff --git a/ggml/src/ggml-cuda/mean.cu b/ggml/src/ggml-cuda/mean.cu index a8f6046e4..64ad7e1d5 100644 --- a/ggml/src/ggml-cuda/mean.cu +++ b/ggml/src/ggml-cuda/mean.cu @@ -18,7 +18,7 @@ void ggml_cuda_op_mean(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { GGML_ASSERT(src0->type == GGML_TYPE_F32); GGML_ASSERT(dst->type == GGML_TYPE_F32); - GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(ggml_is_contiguous_rows(src0)); const int64_t ncols = src0->ne[0]; const int64_t nrows = ggml_nrows(src0); @@ -65,13 +65,20 @@ void ggml_cuda_op_mean(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { // Heuristic for block size selection to optimize occupancy. // See discussion in: https://github.com/ggml-org/llama.cpp/pull/15132 + dim3 block_dims; if ((nrows / nsm) < 2) { - const dim3 block_dims(512, 1, 1); - const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream); - ggml_cuda_kernel_launch(reduce_rows_f32, launch_params, src0_d, dst_d, ncols); + block_dims = dim3(512, 1, 1); } else { - const dim3 block_dims(ncols < 1024 ? 32 : 128, 1, 1); - const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream); - ggml_cuda_kernel_launch(reduce_rows_f32, launch_params, src0_d, dst_d, ncols); + block_dims = dim3(ncols < 1024 ? 32 : 128, 1, 1); } + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream); + + if (ggml_is_contiguous(src0)) { + ggml_cuda_kernel_launch(reduce_rows_f32, launch_params, src0_d, dst_d, ncols); + return; + } + + const char * src0_d_bytes = (const char *) src0->data; + ggml_cuda_kernel_launch(reduce_rows_f32_strided, launch_params, src0_d_bytes, dst_d, ncols, + src0->ne[1], src0->ne[2], src0->nb[1], src0->nb[2], src0->nb[3]); } diff --git a/ggml/src/ggml-cuda/reduce_rows.cuh b/ggml/src/ggml-cuda/reduce_rows.cuh index 968c47aa2..111fd838a 100644 --- a/ggml/src/ggml-cuda/reduce_rows.cuh +++ b/ggml/src/ggml-cuda/reduce_rows.cuh @@ -1,11 +1,6 @@ #include "common.cuh" -// Row reduction kernel template - compute sum (norm=false) or mean (norm=true) -template -static __global__ void reduce_rows_f32(const float * x_ptr, float * dst_ptr, const int ncols) { - const float * GGML_CUDA_RESTRICT x = x_ptr; - float * GGML_CUDA_RESTRICT dst = dst_ptr; - const int row = blockIdx.x; +static __device__ __forceinline__ float reduce_row_f32(const float * x, const int ncols) { const int col = threadIdx.x; float sum = 0.0f; @@ -17,7 +12,7 @@ static __global__ void reduce_rows_f32(const float * x_ptr, float * dst_ptr, con for (int i = col; i < ncols;) { for (int j = 0; j < num_unroll; ++j) { if (i < ncols) { - temp[j] = x[row * ncols + i]; + temp[j] = x[i]; } else { temp[j] = 0; } @@ -35,6 +30,40 @@ static __global__ void reduce_rows_f32(const float * x_ptr, float * dst_ptr, con __shared__ float shared_vals[32]; sum = block_reduce(sum, shared_vals); + return sum; +} + +// Row reduction kernel template - compute sum (norm=false) or mean (norm=true) +template +static __global__ void reduce_rows_f32(const float * x_ptr, float * dst_ptr, const int ncols) { + float * GGML_CUDA_RESTRICT dst = dst_ptr; + const int64_t row = blockIdx.x; + const int col = threadIdx.x; + + const float * GGML_CUDA_RESTRICT x = x_ptr + row*ncols; + const float sum = reduce_row_f32(x, ncols); + + if (col != 0) { + return; + } + + dst[row] = norm ? sum / ncols : sum; +} + +template +static __global__ void reduce_rows_f32_strided(const char * x_ptr, float * dst_ptr, const int ncols, + const int64_t ne1, const int64_t ne2, const int64_t nb1, const int64_t nb2, const int64_t nb3) { + float * GGML_CUDA_RESTRICT dst = dst_ptr; + const int64_t row = blockIdx.x; + const int col = threadIdx.x; + + const int64_t i1 = row % ne1; + const int64_t i2 = (row / ne1) % ne2; + const int64_t i3 = row / (ne1 * ne2); + + const float * GGML_CUDA_RESTRICT x = (const float *) (x_ptr + i1*nb1 + i2*nb2 + i3*nb3); + const float sum = reduce_row_f32(x, ncols); + if (col != 0) { return; } diff --git a/ggml/src/ggml-cuda/sumrows.cu b/ggml/src/ggml-cuda/sumrows.cu index 0003658ca..aa8342b5f 100644 --- a/ggml/src/ggml-cuda/sumrows.cu +++ b/ggml/src/ggml-cuda/sumrows.cu @@ -24,24 +24,30 @@ void ggml_cuda_op_sum_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { GGML_ASSERT(src0->type == GGML_TYPE_F32); GGML_ASSERT( dst->type == GGML_TYPE_F32); - GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(ggml_is_contiguous_rows(src0)); const int64_t ncols = src0->ne[0]; const int64_t nrows = ggml_nrows(src0); + if (ggml_is_contiguous(src0)) { + sum_rows_f32_cuda(src0_d, dst_d, ncols, nrows, stream); + return; + } + const dim3 block_nums(nrows, 1, 1); const int id = ggml_cuda_get_device(); const int nsm = ggml_cuda_info().devices[id].nsm; + dim3 block_dims; if ((nrows / nsm) < 2) { // Increase num threads to 512 for small nrows to better hide the latency - const dim3 block_dims(512, 1, 1); - const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream); - ggml_cuda_kernel_launch(reduce_rows_f32, launch_params, src0_d, dst_d, ncols); + block_dims = dim3(512, 1, 1); } else { // Enough active SMs to hide latency, use smaller blocks to allow better scheduling - const dim3 block_dims(ncols < 1024 ? 32 : 128, 1, 1); - const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream); - ggml_cuda_kernel_launch(reduce_rows_f32, launch_params, src0_d, dst_d, ncols); + block_dims = dim3(ncols < 1024 ? 32 : 128, 1, 1); } + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream); + const char * src0_d_bytes = (const char *) src0->data; + ggml_cuda_kernel_launch(reduce_rows_f32_strided, launch_params, src0_d_bytes, dst_d, ncols, + src0->ne[1], src0->ne[2], src0->nb[1], src0->nb[2], src0->nb[3]); }