Public/koboldcpp
1
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-09-19 17:25:07 +02:00
Code Issues Packages Projects Releases Wiki Activity

cuda: support row-contiguous SUM_ROWS (#26308)

Browse Source 
* 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
This commit is contained in:
Mohamed Elashri
2026-09-15 12:39:29 +02:00
committed by GitHub
parent 9e71716247
commit 5431581326
5 changed files with 83 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cuda/ggml-cuda.cu
+2
View File
@@ -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:
ggml/src/ggml-cuda/mean.cu
+14 -7
View File
@@ -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);
ggml_cuda_kernel_launch(reduce_rows_f32</*norm=*/true>, 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</*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]);
}
ggml/src/ggml-cuda/reduce_rows.cuh
+36 -7
View File
@@ -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;
}
ggml/src/ggml-cuda/sumrows.cu
+13 -7
View File
@@ -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]);
}
tests/test-backend-ops.cpp
+18 -3
View File
@@ -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 }));