nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

New op: ggml_sum_rows_ext (#2132) 

* Add ggml_sum_rows_ext

* openPangu: use ggml_sum_rows_ext also in mhc_post

* openPangu: use ggml_sum_rows_ext also in mhc_tail

* Minor
This commit is contained in:
Kawrakow 2026-07-15 16:21:17 +03:00 committed by GitHub
parent 5596a41a38
commit 1fddd12ba8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 205 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
ggml/include/ggml.h
View File

@ -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,

77
ggml/src/ggml-cuda/sumrows.cu
View File

@ -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));

134
ggml/src/ggml.c
View File

@ -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(
@ -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(

27
src/graphs/build_openpangu.cpp
View File

@ -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));
@ -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) {
@ -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.
@ -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]
@ -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).