diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index d3c540da..38be6e64 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -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, diff --git a/ggml/src/ggml-cuda/sumrows.cu b/ggml/src/ggml-cuda/sumrows.cu index c4929e81..eab05305 100644 --- a/ggml/src/ggml-cuda/sumrows.cu +++ b/ggml/src/ggml-cuda/sumrows.cu @@ -75,10 +75,87 @@ static void sum_rows_div_f32_cuda(const float * x, float * dst, const int ncols, k_sum_rows_div_f32<<>>(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<<>>( + (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<<>>( + (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<<>>( + (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)); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 9ec3886c..eff90454 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -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( diff --git a/src/graphs/build_openpangu.cpp b/src/graphs/build_openpangu.cpp index 37371f13..b0d9ca20 100644 --- a/src/graphs/build_openpangu.cpp +++ b/src/graphs/build_openpangu.cpp @@ -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).