Skip to content

New op: ggml_sum_rows_ext #2132

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ikawrakow merged 4 commits into from
Jul 15, 2026
+205 −37
Merged

New op: ggml_sum_rows_ext #2132

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ae824ef
Add ggml_sum_rows_ext
ikawrakow Jul 14, 2026
7d71d6e
openPangu: use ggml_sum_rows_ext also in mhc_post
ikawrakow Jul 14, 2026
e3e6c11
openPangu: use ggml_sum_rows_ext also in mhc_tail
ikawrakow Jul 14, 2026
d477e24
Minor
ikawrakow Jul 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions ggml/include/ggml.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1235,6 +1235,10 @@ extern "C" {
GGML_API struct ggml_tensor * ggml_sum_rows(
struct ggml_context * ctx,
struct ggml_tensor * a);
GGML_API struct ggml_tensor * ggml_sum_rows_ext(
struct ggml_context * ctx,
struct ggml_tensor * a,
int dim);

GGML_API struct ggml_tensor * ggml_cumsum(
struct ggml_context * ctx,
Expand Down
77 changes: 77 additions & 0 deletions ggml/src/ggml-cuda/sumrows.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,10 +75,87 @@ static void sum_rows_div_f32_cuda(const float * x, float * dst, const int ncols,
k_sum_rows_div_f32<<<block_nums, block_dims, 0, stream>>>(x, dst, ncols, s, b);
}

static __global__ void k_sum_rows_any_f32(const char * x, char * y, const int ncols,
const int nsum, const int ne1,
size_t nbsum, size_t nb01, size_t nb02, size_t nb1, size_t nb2) {

int col = blockIdx.y*blockDim.x + threadIdx.x;
if (col >= ncols) {
return;
}
int row = blockIdx.x;
int i2 = row/ne1;
int i1 = row%ne1;

//x += i1*nb01 + i2*nb02 + col*sizeof(float);
//y += i1*nb1 + i2*nb2 + col*sizeof(float);

//float sum = 0.0f;
//for (int is = 0; is < nsum; ++is) {
// sum += *(const float *)x;
// x += nbsum;
//}

//*y = sum;

const float * xf = (const float *)(x + i1*nb01 + i2*nb02);
float * yf = (float *)(y + i1*nb1 + i2*nb2);

float sum = 0.0f;
for (int is = 0; is < nsum; ++is) {
sum += xf[is*nbsum + col];
}

yf[col] = sum;
}

void ggml_cuda_op_sum_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const ggml_tensor * src0 = dst->src[0];
cudaStream_t stream = ctx.stream();

int dim = dst->op_params[0];
if (dim > 0) {
//printf("I'm here (%d): %ld x %ld x %ld x %ld -> %ld x %ld x %ld x %ld\n", dim,
// src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
GGML_ASSERT(src0->nb[0] == sizeof(float) && dst->nb[0] == sizeof(float));
constexpr int kBlockSize = 256;
int nblock = (dst->ne[0] + kBlockSize - 1)/kBlockSize;
// Note: we use dim y as the number of blocks because we do not expect
// tensors that may overflow the 65536 limit for this dim, i.e.,
// we do not expect tensors to have row sizes > 16777216.
if (dim == 1) {
GGML_ASSERT(dst->ne[1] == 1);
GGML_ASSERT(src0->ne[0] == dst->ne[0] && src0->ne[2] == dst->ne[2] && src0->ne[3] == dst->ne[3]);
int nrows = dst->ne[2]*dst->ne[3];
dim3 grid(nrows, nblock, 1);
k_sum_rows_any_f32<<<grid, kBlockSize, 0, ctx.stream()>>>(
(const char *)src0->data, (char *)dst->data, dst->ne[0],
src0->ne[1], src0->ne[2],
src0->nb[1]/sizeof(float), src0->nb[2], src0->nb[3], dst->nb[2], dst->nb[3]);
//src0->nb[1], src0->nb[2], src0->nb[3], dst->nb[2], dst->nb[3]);
}
else if (dim == 2) {
int nrows = dst->ne[1]*dst->ne[3];
dim3 grid(nrows, nblock, 1);
k_sum_rows_any_f32<<<grid, kBlockSize, 0, ctx.stream()>>>(
(const char *)src0->data, (char *)dst->data, dst->ne[0],
src0->ne[2], src0->ne[1],
src0->nb[2]/sizeof(float), src0->nb[1], src0->nb[3], dst->nb[1], dst->nb[3]);
}
else if (dim == 3) {
int nrows = dst->ne[1]*dst->ne[2];
dim3 grid(nrows, nblock, 1);
k_sum_rows_any_f32<<<grid, kBlockSize, 0, ctx.stream()>>>(
(const char *)src0->data, (char *)dst->data, dst->ne[0],
src0->ne[3], src0->ne[1],
src0->nb[3]/sizeof(float), src0->nb[1], src0->nb[2], dst->nb[1], dst->nb[2]);
}
else {
GGML_ABORT("Unsupported sum_rows dimension");
}
return;
}

GGML_ASSERT(src0->type == GGML_TYPE_F32);
GGML_ASSERT( dst->type == GGML_TYPE_F32);
GGML_ASSERT(ggml_is_contiguous(src0));
Expand Down
134 changes: 110 additions & 24 deletions ggml/src/ggml.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6789,6 +6789,26 @@ struct ggml_tensor * ggml_sum_rows(
return result;
}

struct ggml_tensor * ggml_sum_rows_ext(
struct ggml_context * ctx,
struct ggml_tensor * a,
int dim) {
GGML_ASSERT(dim >= 0 && dim < GGML_MAX_DIMS);
if (dim == 0) return ggml_sum_rows(ctx, a);

int64_t ne[GGML_MAX_DIMS];
for (int i = 0; i < GGML_MAX_DIMS; ++i) ne[i] = a->ne[i];
ne[dim] = 1;

struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, ne);

result->op = GGML_OP_SUM_ROWS;
result->src[0] = a;
result->op_params[0] = dim;

return result;
}

// ggml_cumsum

struct ggml_tensor * ggml_cumsum(
Expand Down Expand Up @@ -14501,37 +14521,103 @@ static void ggml_compute_forward_sum_rows_f32(

GGML_TENSOR_UNARY_OP_LOCALS

GGML_ASSERT(ne0 == 1);
GGML_ASSERT(ne1 == ne01);
GGML_ASSERT(ne2 == ne02);
GGML_ASSERT(ne3 == ne03);

int ith = params->ith;
int nth = params->nth;

//if (params->ith == 0) printf("%s(%s): %ld x %ld x %ld x %ld\n", __func__, dst->name, ne00, ne1, ne2, ne3);
int dim = dst->op_params[0];

int nrows = ggml_nrows(src0);
int nrows_per_thread = (nrows + nth - 1)/nth;
int first_row = nrows_per_thread*ith;
int last_row = MIN(first_row + nrows_per_thread, nrows);
if (dim == 0) {
GGML_ASSERT(ne0 == 1);
GGML_ASSERT(ne1 == ne01);
GGML_ASSERT(ne2 == ne02);
GGML_ASSERT(ne3 == ne03);
int nrows = ggml_nrows(src0);
int nrows_per_thread = (nrows + nth - 1)/nth;
int first_row = nrows_per_thread*ith;
int last_row = MIN(first_row + nrows_per_thread, nrows);

for (int ir = first_row; ir < last_row; ++ir) {
int i3 = ir / (ne01*ne02);
int i2 = (ir - i3*ne01*ne02)/ne01;
int i1 = ir - i3*ne01*ne02 - i2*ne01;
const float * src_row = (const float *)((const char *)src0->data + i1*nb01 + i2*nb02 + i3*nb03);
float * dst_row = ( float *)(( char *)dst->data + i1*nb1 + i2*nb2 + i3*nb3);
float row_sum = 0;
ggml_vec_sum_f32(ne00, &row_sum, src_row);
if (!isfinite(row_sum)) {
fprintf(stderr, "Oops(%s, %s): found %g for i1 = %d, i2 = %d, i3 = %d. ne00 = %d\n", __func__, dst->name,
(double)row_sum, (int)i1, (int)i2, (int)i3, (int)ne00);
GGML_ABORT("Fatal error");
for (int ir = first_row; ir < last_row; ++ir) {
int i3 = ir / (ne01*ne02);
int i2 = (ir - i3*ne01*ne02)/ne01;
int i1 = ir - i3*ne01*ne02 - i2*ne01;
const float * src_row = (const float *)((const char *)src0->data + i1*nb01 + i2*nb02 + i3*nb03);
float * dst_row = ( float *)(( char *)dst->data + i1*nb1 + i2*nb2 + i3*nb3);
float row_sum = 0;
ggml_vec_sum_f32(ne00, &row_sum, src_row);
if (!isfinite(row_sum)) {
fprintf(stderr, "Oops(%s, %s): found %g for i1 = %d, i2 = %d, i3 = %d. ne00 = %d\n", __func__, dst->name,
(double)row_sum, (int)i1, (int)i2, (int)i3, (int)ne00);
GGML_ABORT("Fatal error");
}
dst_row[0] = row_sum;
}
} else {
for (int i = 0; i < GGML_MAX_DIMS; ++i) {
if (i == dim) {
GGML_ASSERT(dst->ne[i] == 1);
} else {
GGML_ASSERT(dst->ne[i] == src0->ne[i]);
}
}
int n = dst->ne[0];
if (dim == 1) {
int nrows = src0->ne[2]*src0->ne[3];
int nrows_per_thread = (nrows + nth - 1)/nth;
int first_row = nrows_per_thread*ith;
int last_row = MIN(first_row + nrows_per_thread, nrows);
for (int ir = first_row; ir < last_row; ++ir) {
int i3 = ir/src0->ne[2];
int i2 = ir - i3*src0->ne[2];
const char * csrc = (const char *)src0->data + i2*nb02 + i3*nb03;
char * cdst = (char *)dst->data + i2*nb2 + i3*nb3;
memcpy(cdst, csrc, n*sizeof(float));
float * y = (float *)cdst;
for (int i1 = 1; i1 < src0->ne[1]; ++i1) {
csrc += src0->nb[1];
const float * x = (const float *)csrc;
for (int j = 0; j < n; ++j) y[j] += x[j];
}
}
}
else if (dim == 2) {
int nrows = src0->ne[1]*src0->ne[3];
int nrows_per_thread = (nrows + nth - 1)/nth;
int first_row = nrows_per_thread*ith;
int last_row = MIN(first_row + nrows_per_thread, nrows);
for (int ir = first_row; ir < last_row; ++ir) {
int i3 = ir/src0->ne[1];
int i1 = ir - i3*src0->ne[1];
const char * csrc = (const char *)src0->data + i1*nb01 + i3*nb03;
char * cdst = (char *)dst->data + i1*nb1 + i3*nb3;
memcpy(cdst, csrc, n*sizeof(float));
float * y = (float *)cdst;
for (int i2 = 1; i2 < src0->ne[2]; ++i2) {
csrc += src0->nb[2];
const float * x = (const float *)csrc;
for (int j = 0; j < n; ++j) y[j] += x[j];
}
}
}
else {
int nrows = src0->ne[1]*src0->ne[2];
int nrows_per_thread = (nrows + nth - 1)/nth;
int first_row = nrows_per_thread*ith;
int last_row = MIN(first_row + nrows_per_thread, nrows);
for (int ir = first_row; ir < last_row; ++ir) {
int i2 = ir/src0->ne[1];
int i1 = ir - i2*src0->ne[1];
const char * csrc = (const char *)src0->data + i1*nb01 + i2*nb02;
char * cdst = (char *)dst->data + i1*nb1 + i2*nb2;
memcpy(cdst, csrc, n*sizeof(float));
float * y = (float *)cdst;
for (int i3 = 1; i3 < src0->ne[3]; ++i3) {
csrc += src0->nb[3];
const float * x = (const float *)csrc;
for (int j = 0; j < n; ++j) y[j] += x[j];
}
}
}
dst_row[0] = row_sum;
}

}

static void ggml_compute_forward_sum_rows(
Expand Down
27 changes: 14 additions & 13 deletions src/graphs/build_openpangu.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1028,9 +1028,11 @@ ggml_cgraph * llm_build_context::build_openpangu() {
auto mhc_pre = [&](ggml_tensor * Rin, ggml_tensor * phi, ggml_tensor * alpha,
ggml_tensor * beta, ggml_tensor * gamma,
ggml_tensor ** h_post_out, ggml_tensor ** h_res_out) {
ggml_tensor * flat = ggml_reshape_2d(ctx0, ggml_cont(ctx0, Rin), n_embd * S, n_tokens);
ggml_tensor * normed = ggml_rms_norm(ctx0, flat, hparams.f_norm_rms_eps);
normed = ggml_mul(ctx0, normed, gamma); // [S*H, T]
if (!ggml_is_contiguous(Rin)) {
Rin = ggml_cont(ctx0, Rin);
}
ggml_tensor * flat = ggml_reshape_2d(ctx0, Rin, n_embd * S, n_tokens);
ggml_tensor * normed = ggml_fused_rms_norm(ctx0, flat, gamma, hparams.f_norm_rms_eps);
ggml_tensor * mixes = ggml_mul_mat(ctx0, phi, normed); // [(S+2)*S, T]
ggml_tensor * h_pre = ggml_view_2d(ctx0, mixes, S, n_tokens, mixes->nb[1], 0);
ggml_tensor * h_post = ggml_view_2d(ctx0, mixes, S, n_tokens, mixes->nb[1], S*ggml_element_size(mixes));
Expand All @@ -1045,16 +1047,19 @@ ggml_cgraph * llm_build_context::build_openpangu() {
h_pre = ggml_sigmoid(ctx0, h_pre); // [S,T] (+eps omitted, inert)

// combine: x[h,t] = sum_s h_pre[s,t] * R[h,s,t]
ggml_tensor * hpre3 = ggml_reshape_3d(ctx0, ggml_cont(ctx0, h_pre), 1, S, n_tokens);
ggml_tensor * hpre3 = ggml_reshape_3d(ctx0, h_pre, 1, S, n_tokens);
ggml_tensor * weighted = ggml_mul(ctx0, Rin, hpre3); // [H,S,T]
ggml_tensor * wperm = ggml_cont(ctx0, ggml_permute(ctx0, weighted, 1, 0, 2, 3)); // [S,H,T]
ggml_tensor * x = ggml_reshape_2d(ctx0, ggml_sum_rows(ctx0, wperm), n_embd, n_tokens); // sum over S
ggml_tensor * x = ggml_reshape_2d(ctx0, ggml_sum_rows_ext(ctx0, weighted, 1), n_embd, n_tokens);
ggml_build_forward_expand(gf, x);

*h_post_out = ggml_cont(ctx0, h_post);
*h_res_out = ggml_cont(ctx0, h_res);
return x;
};

ggml_tensor repeater;
repeater.ne[0] = n_embd; repeater.ne[1] = S; repeater.ne[2] = n_tokens; repeater.ne[3] = 1;

// mHC post: R_new[h,s,t] = h_post[s,t]*y[h,t] + sum_j m[s,j,t]*R[h,j,t]
auto mhc_post = [&](ggml_tensor * y, ggml_tensor * h_post, ggml_tensor * Rin,
ggml_tensor * alpha, ggml_tensor * beta, ggml_tensor * h_res) {
Expand All @@ -1072,8 +1077,7 @@ ggml_cgraph * llm_build_context::build_openpangu() {
// term1: h_post[s,t]*y[h,t] -> [H,S,T]
ggml_tensor * y3 = ggml_reshape_3d(ctx0, y, n_embd, 1, n_tokens);
ggml_tensor * hpost3 = ggml_reshape_3d(ctx0, ggml_cont(ctx0, h_post), 1, S, n_tokens);
ggml_tensor * term1 = ggml_mul(ctx0, ggml_repeat(ctx0, y3,
ggml_new_tensor_3d(ctx0, y->type, n_embd, S, n_tokens)), hpost3);
ggml_tensor * term1 = ggml_mul(ctx0, ggml_repeat(ctx0, y3, &repeater), hpost3);

// term2: sum_j m[s,j,t]*R[h,j,t]. For each out-stream s, weight over input streams j.
// Build via: for stream axis, matmul R[H, j, t] with m[j, s, t] batched over t.
Expand All @@ -1085,9 +1089,7 @@ ggml_cgraph * llm_build_context::build_openpangu() {
ggml_tensor * m_s = ggml_cont(ctx0, ggml_view_2d(ctx0, m, S, n_tokens, m->nb[2], s*m->nb[1]));
ggml_tensor * m_s3 = ggml_reshape_3d(ctx0, m_s, 1, S, n_tokens); // [1,S,T]
ggml_tensor * acc = ggml_mul(ctx0, Rin, m_s3); // [H,S,T]
ggml_tensor * accp = ggml_cont(ctx0, ggml_permute(ctx0, acc, 1, 0, 2, 3)); // [S,H,T]
ggml_tensor * summed = ggml_reshape_2d(ctx0, ggml_sum_rows(ctx0, accp), n_embd, n_tokens); // [H,T]
summed = ggml_reshape_3d(ctx0, summed, n_embd, 1, n_tokens);
ggml_tensor * summed = ggml_sum_rows_ext(ctx0, acc, 1); // [H,1,T]
term2 = term2 ? ggml_concat(ctx0, term2, summed, 1) : summed; // -> [H,S,T]
}
return ggml_add(ctx0, term1, term2); // [H,S,T]
Expand Down Expand Up @@ -1167,8 +1169,7 @@ ggml_cgraph * llm_build_context::build_openpangu() {
w = ggml_sigmoid(ctx0, ggml_add(ctx0, ggml_mul(ctx0, w, a_pre), model.mhc_merge_beta)); // [S,T]
ggml_tensor * w3 = ggml_reshape_3d(ctx0, ggml_cont(ctx0, w), 1, S, n_tokens);
ggml_tensor * weighted = ggml_mul(ctx0, R, w3); // [H,S,T]
ggml_tensor * wperm = ggml_cont(ctx0, ggml_permute(ctx0, weighted, 1, 0, 2, 3)); // [S,H,T]
cur = ggml_reshape_2d(ctx0, ggml_sum_rows(ctx0, wperm), n_embd, n_tokens);
cur = ggml_reshape_2d(ctx0, ggml_sum_rows_ext(ctx0, weighted, 1), n_embd, n_tokens);
}

// select only the output tokens (the framework binds n_outputs rows, not all n_tokens).
Expand Down