Commit 543158132 for llama.cpp
commit 54315813269112dd0baed7112ec87ad93a8218ca
Author: Mohamed Elashri <git@elashri.com>
Date: Tue Sep 15 12:39:29 2026 +0200
cuda: support row-contiguous SUM_ROWS (#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
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</*norm=*/true>, 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);
+ 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</*norm=*/true>, 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</*norm=*/true>, 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 <bool norm>
-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<block_reduce_method::SUM>(sum, shared_vals);
+ return sum;
+}
+
+// Row reduction kernel template - compute sum (norm=false) or mean (norm=true)
+template <bool norm>
+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 <bool norm>
+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</*norm=*/false>, 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</*norm=*/false>, 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</*norm=*/false>, 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/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 0e074770d..1616004e0 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -7126,20 +7126,32 @@ struct test_sum_rows : public test_case {
struct test_mean : public test_case {
const ggml_type type;
const std::array<int64_t, 4> ne;
+ const bool permute;
+ const bool slice;
std::string vars() override {
- return VARS_TO_STR2(type, ne);
+ return VARS_TO_STR4(type, ne, permute, slice);
}
test_mean(ggml_type type = GGML_TYPE_F32,
- std::array<int64_t, 4> ne = {10, 5, 4, 3})
- : type(type), ne(ne) {}
+ std::array<int64_t, 4> ne = {10, 5, 4, 3},
+ bool permute = false, bool slice = false)
+ : type(type), ne(ne), permute(permute), slice(slice) {}
ggml_tensor * build_graph(ggml_context * ctx) override {
ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data());
ggml_set_param(a);
ggml_set_name(a, "a");
+ if (slice) {
+ a = ggml_view_4d(ctx, a,
+ ne[0], ne[1], ne[2] / 2, ne[3] - 1,
+ a->nb[1], a->nb[2] * 2, a->nb[3], /*offset=*/a->nb[3]);
+ }
+ if (permute) {
+ a = ggml_permute(ctx, a, 0, 2, 3, 1);
+ }
+
ggml_tensor * out = ggml_mean(ctx, a);
ggml_set_name(out, "out");
@@ -10470,6 +10482,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_mean(GGML_TYPE_F32, { 32, 1, 1, 1 }));
test_cases.emplace_back(new test_mean(GGML_TYPE_F32, { 32, 256, 1, 1 }));
test_cases.emplace_back(new test_mean(GGML_TYPE_F32, { 32768, 1, 1, 1 }));
+ test_cases.emplace_back(new test_mean(GGML_TYPE_F32, { 11, 5, 6, 3 }, true, false));
+ test_cases.emplace_back(new test_mean(GGML_TYPE_F32, { 11, 5, 6, 3 }, false, true));
+ test_cases.emplace_back(new test_mean(GGML_TYPE_F32, { 11, 5, 6, 3 }, true, true));
test_cases.emplace_back(new test_sum(GGML_TYPE_F32, { 33, 1, 1, 1 }));
test_cases.emplace_back(new test_sum(GGML_TYPE_F32, { 33, 1024, 1, 1 }));
test_cases.emplace_back(new test_sum(GGML_TYPE_F32, { 33, 256, 1, 1 }));