diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 94e09840..2854614e 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -699,6 +699,7 @@ extern "C" { GGML_OP_REDUCE, GGML_OP_FAKE_CPY, GGML_OP_FUSED_NORM, + GGML_OP_FUSED_RMS_RMS_ADD, GGML_OP_COUNT, }; @@ -1571,6 +1572,14 @@ extern "C" { struct ggml_tensor * a, float eps); + GGML_API struct ggml_tensor * ggml_fused_rms_rms_add( + struct ggml_context * ctx, + struct ggml_tensor * x1, + struct ggml_tensor * c1, + struct ggml_tensor * x2, + struct ggml_tensor * c2, + float eps); + // a - x // b - dy GGML_API struct ggml_tensor * ggml_rms_norm_back( diff --git a/ggml/src/ggml-cuda.cu b/ggml/src/ggml-cuda.cu index 4d1ebd5c..d1abfc13 100644 --- a/ggml/src/ggml-cuda.cu +++ b/ggml/src/ggml-cuda.cu @@ -3665,6 +3665,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg ggml_cuda_op_fused_rms_norm(ctx, dst); } break; + case GGML_OP_FUSED_RMS_RMS_ADD: + ggml_cuda_op_fused_rms_rms_add(ctx, dst); + break; case GGML_OP_FUSED_NORM: ggml_cuda_op_fused_rms_norm(ctx, dst, true); break; @@ -4658,6 +4661,7 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons case GGML_OP_DIV: case GGML_OP_SUB: case GGML_OP_FUSED_RMS_NORM: + case GGML_OP_FUSED_RMS_RMS_ADD: case GGML_OP_SCALE: case GGML_OP_SOFTCAP: case GGML_OP_SQR: diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index 61d90a0c..ebc1cce7 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -951,3 +951,104 @@ void ggml_cuda_op_fused_rms_rms_norm([[maybe_unused]] ggml_backend_cuda_context } + +template +static __global__ void fused_rms_rms_add_f32(int ncols, int nrows, float * dst, + const src_t * x1, const float * c1, const src_t * x2, const float * c2, float eps) { + const int row = blockIdx.x*blockDim.y + threadIdx.y; + const int tid = threadIdx.x; + + auto x1_row = x1 + row*ncols; + auto x2_row = x2 + row*ncols; + + float tmp1 = 0.0f, tmp2 = 0.0f; + + for (int col = tid; col < ncols; col += block_size) { + const float xi1 = (float)x1_row[col]; + const float xi2 = (float)x2_row[col]; + tmp1 += xi1 * xi1; + tmp2 += xi2 * xi2; + } + + tmp1 = warp_reduce_sum(tmp1); + tmp2 = warp_reduce_sum(tmp2); + if (block_size > WARP_SIZE) { + __shared__ float s_sum[2*WARP_SIZE]; + int warp_id = threadIdx.x / WARP_SIZE; + int lane_id = threadIdx.x % WARP_SIZE; + if (lane_id == 0) { + s_sum[2*warp_id+0] = tmp1; + s_sum[2*warp_id+1] = tmp2; + } + __syncthreads(); + tmp1 = lane_id < block_size/WARP_SIZE ? s_sum[2*lane_id+0] : 0.0f; + tmp2 = lane_id < block_size/WARP_SIZE ? s_sum[2*lane_id+1] : 0.0f; + tmp1 = warp_reduce_sum(tmp1); + tmp2 = warp_reduce_sum(tmp2); + } + + const float mean1 = tmp1 / ncols; + const float mean2 = tmp2 / ncols; + const float scale1 = rsqrtf(mean1 + eps); + const float scale2 = rsqrtf(mean2 + eps); + + dst += row*ncols; + + for (int col = tid; col < ncols; col += block_size) { + dst[col] = scale1 * c1[col] * (float)x1_row[col] + scale2 * c2[col] * (float)x2_row[col]; + } +} + +template +static void fused_rms_rms_add_f32_cuda(int ncols, int nrows, float * dst, + const src_t * x1, const float * c1, const src_t * x2, const float * c2, + float eps, cudaStream_t stream) { + if (ncols < 1024) { + const dim3 block_dims(256, 1, 1); + fused_rms_rms_add_f32<256><<>>(ncols, nrows, dst, x1, c1, x2, c2, eps); + } else { + const dim3 block_dims(1024, 1, 1); + fused_rms_rms_add_f32<1024><<>>(ncols, nrows, dst, x1, c1, x2, c2, eps); + } +} + +void ggml_cuda_op_fused_rms_rms_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + GGML_ASSERT(ggml_are_same_shape(dst->src[0], dst->src[2])); + GGML_ASSERT(ggml_are_same_shape(dst->src[0], dst)); + GGML_ASSERT(ggml_is_contiguous(dst->src[0])); + GGML_ASSERT(ggml_is_contiguous(dst->src[2])); + GGML_ASSERT(ggml_is_contiguous(dst)); + GGML_ASSERT(ggml_nrows(dst->src[1]) == 1 && dst->src[1]->ne[0] == dst->src[0]->ne[0]); + GGML_ASSERT(ggml_nrows(dst->src[3]) == 1 && dst->src[3]->ne[0] == dst->src[2]->ne[0]); + GGML_ASSERT(dst->src[0]->type == dst->src[2]->type); + GGML_ASSERT(dst->src[1]->type == GGML_TYPE_F32 && dst->src[3]->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + float eps; + memcpy(&eps, dst->op_params, sizeof(float)); + + int nrows = ggml_nrows(dst); + int ncols = dst->ne[0]; + + if (dst->src[0]->type == GGML_TYPE_F32) { + fused_rms_rms_add_f32_cuda(ncols, nrows, (float *)dst->data, + (const float *)dst->src[0]->data, (const float *)dst->src[1]->data, + (const float *)dst->src[2]->data, (const float *)dst->src[3]->data, + eps, ctx.stream()); + } + else if (dst->src[0]->type == GGML_TYPE_F16) { + fused_rms_rms_add_f32_cuda(ncols, nrows, (float *)dst->data, + (const half *)dst->src[0]->data, (const float *)dst->src[1]->data, + (const half *)dst->src[2]->data, (const float *)dst->src[3]->data, + eps, ctx.stream()); + } + else if (dst->src[0]->type == GGML_TYPE_BF16) { + fused_rms_rms_add_f32_cuda(ncols, nrows, (float *)dst->data, + (const nv_bfloat16 *)dst->src[0]->data, (const float *)dst->src[1]->data, + (const nv_bfloat16 *)dst->src[2]->data, (const float *)dst->src[2]->data, + eps, ctx.stream()); + } + else { + GGML_ABORT("Not implemented"); + } +} diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index 0c6c16a6..513af1e2 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -15,3 +15,5 @@ void ggml_cuda_op_fused_add_rms_norm(ggml_backend_cuda_context & ctx, ggml_tenso void ggml_cuda_op_fused_add_add_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * add1, ggml_tensor * add2, ggml_tensor * dst); void ggml_cuda_op_fused_rms_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * rms1, ggml_tensor * rms2); + +void ggml_cuda_op_fused_rms_rms_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 8e81d633..bc049720 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -4319,9 +4319,10 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "REDUCE", "FAKE_CPY", "FUSED_NORM", + "FUSED_RMS_RMS_ADD", }; -static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); +static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -4438,9 +4439,11 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "reduce(x1,x2,...)", "fake_cpy(x,y)", "norm(x,y)", + "rms(x1)+rms(x2)", + }; -static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); +static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -7474,6 +7477,36 @@ struct ggml_tensor * ggml_rms_norm_inplace( return ggml_rms_norm_impl(ctx, a, eps, true); } +struct ggml_tensor * ggml_fused_rms_rms_add( + struct ggml_context * ctx, + struct ggml_tensor * x1, + struct ggml_tensor * c1, + struct ggml_tensor * x2, + struct ggml_tensor * c2, + float eps) { + + GGML_ASSERT(ggml_is_contiguous(x1) && ggml_is_contiguous(x2)); + GGML_ASSERT(ggml_are_same_shape(x1, x2)); + GGML_ASSERT(ggml_nrows(c1) == 1 && ggml_nrows(c2) == 1); + GGML_ASSERT(x1->ne[0] == c1->ne[0] && x2->ne[0] == c2->ne[0]); + GGML_ASSERT(x1->type == x2->type); + GGML_ASSERT(x1->type == GGML_TYPE_F16 || x1->type == GGML_TYPE_BF16 || x1->type == GGML_TYPE_F32); + GGML_ASSERT(c1->type == GGML_TYPE_F32 && c2->type == GGML_TYPE_F32); + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, x1->ne[0], x1->ne[1], x1->ne[2], x1->ne[3]); + + memcpy(result->op_params, &eps, sizeof(eps)); + + result->op = GGML_OP_FUSED_RMS_RMS_ADD; + result->src[0] = x1; + result->src[1] = c1; + result->src[2] = x2; + result->src[3] = c2; + + return result; +} + + static struct ggml_tensor * ggml_fused_rms_norm_impl( struct ggml_context * ctx, struct ggml_tensor * a, @@ -24203,6 +24236,10 @@ static int ggml_compute_forward(struct ggml_compute_params * params, struct ggml { ggml_compute_forward_fused_rms_norm(params, tensor); } break; + case GGML_OP_FUSED_RMS_RMS_ADD: + { + iqk_rms_rms_add(tensor, params->ith, params->nth); + } break; case GGML_OP_FUSED_NORM: { ggml_compute_forward_fused_norm(params, tensor); @@ -25040,6 +25077,7 @@ static void ggml_compute_backward(struct ggml_context * ctx, struct ggml_tensor } } break; case GGML_OP_FUSED_RMS_NORM: + case GGML_OP_FUSED_RMS_RMS_ADD: case GGML_OP_FUSED_NORM: { GGML_ABORT("fatal error"); // TODO: not implemented @@ -26240,6 +26278,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_NORM: case GGML_OP_RMS_NORM: case GGML_OP_FUSED_RMS_NORM: + case GGML_OP_FUSED_RMS_RMS_ADD: case GGML_OP_FUSED_NORM: case GGML_OP_RMS_NORM_BACK: case GGML_OP_GROUP_NORM: diff --git a/ggml/src/iqk/iqk_cpu_ops.cpp b/ggml/src/iqk/iqk_cpu_ops.cpp index 7a91336b..7b87e88d 100644 --- a/ggml/src/iqk/iqk_cpu_ops.cpp +++ b/ggml/src/iqk/iqk_cpu_ops.cpp @@ -823,5 +823,134 @@ bool iqk_ssm_conv4(int nr, int nc, int nt, #else return false; #endif - } +} +namespace { +inline float sum_row_squared(int ncols, const float * x) { + float sum = 0; + int i = 0; +#ifdef __AVX2__ + auto vsum = _mm256_setzero_ps(); + for (; i < ncols - 7; i += 8) { + auto vx = _mm256_loadu_ps(x + i); + vsum = _mm256_fmadd_ps(vx, vx, vsum); + } + sum = hsum_float_8(vsum); +#endif + for (; i < ncols; ++i) sum += x[i]*x[i]; + //for (int j = 0; j < ncols; ++j) sum += x[j]*x[j]; + return sum; +} +inline float sum_row_squared(int ncols, const ggml_half * x) { + float sum = 0; + for (int j = 0; j < ncols; ++j) { + float v = GGML_FP16_TO_FP32(x[j]); + sum += v*v; + } + return sum; +} +inline float sum_row_squared(int ncols, const ggml_bf16_t * x) { + float sum = 0; + for (int j = 0; j < ncols; ++j) { + float v = GGML_BF16_TO_FP32(x[j]); + sum += v*v; + } + return sum; +} +inline void rms_rms_add(int ncols, float scale1, float scale2, const float * x1, const float * x2, const float * c1, const float * c2, float * dst) { + int j = 0; +#ifdef __AVX2__ + auto vs1 = _mm256_set1_ps(scale1); + auto vs2 = _mm256_set1_ps(scale2); + for (; j < ncols - 7; j += 8) { + auto vx1 = _mm256_loadu_ps(x1 + j); + auto vx2 = _mm256_loadu_ps(x2 + j); + auto vc1 = _mm256_loadu_ps(c1 + j); + auto vc2 = _mm256_loadu_ps(c2 + j); + auto vy = _mm256_add_ps(_mm256_mul_ps(_mm256_mul_ps(vs1, vc1), vx1), _mm256_mul_ps(_mm256_mul_ps(vs2, vc2), vx2)); + _mm256_storeu_ps(dst + j, vy); + } +#endif + for (; j < ncols; ++j) { + dst[j] = scale1 * c1[j] * x1[j] + scale2 * c2[j] * x2[j]; + } +} +inline void rms_rms_add(int ncols, float scale1, float scale2, const ggml_half * x1, const ggml_half * x2, const float * c1, const float * c2, float * dst) { + for (int j = 0; j < ncols; ++j) { + float v1 = GGML_FP16_TO_FP32(x1[j]); + float v2 = GGML_FP16_TO_FP32(x2[j]); + dst[j] = scale1 * c1[j] * v1 + scale2 * c2[j] * v2; + } +} +inline void rms_rms_add(int ncols, float scale1, float scale2, const ggml_bf16_t * x1, const ggml_bf16_t * x2, const float * c1, const float * c2, float * dst) { + for (int j = 0; j < ncols; ++j) { + float v1 = GGML_BF16_TO_FP32(x1[j]); + float v2 = GGML_BF16_TO_FP32(x2[j]); + dst[j] = scale1 * c1[j] * v1 + scale2 * c2[j] * v2; + } +} +} + +void iqk_rms_rms_add(struct ggml_tensor * dst, int ith, int nth) { + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + const struct ggml_tensor * src0 = dst->src[0]; + const struct ggml_tensor * src1 = dst->src[1]; + const struct ggml_tensor * src2 = dst->src[2]; + const struct ggml_tensor * src3 = dst->src[3]; + + GGML_ASSERT(ggml_is_contiguous(src0) && ggml_is_contiguous(src2) && ggml_is_contiguous(dst)); + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(ggml_are_same_shape(src2, dst)); + GGML_ASSERT(ggml_nrows(src1) == 1 && ggml_nrows(src3) == 1); + GGML_ASSERT(src0->ne[0] == src1->ne[0] && src2->ne[0] == src3->ne[0]); + GGML_ASSERT(src0->type == src2->type); + GGML_ASSERT(src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16 || src0->type == GGML_TYPE_F32); + + float eps; + memcpy(&eps, dst->op_params, sizeof(float)); + + GGML_ASSERT(eps > 0.0f); + + int nrows = ggml_nrows(dst); + int nrows_per_thread = (nrows + nth - 1)/nth; + int first = ith*nrows_per_thread; + int last = MIN(nrows, first + nrows_per_thread); + + const float * c1 = (float *) src1->data; + const float * c2 = (float *) src3->data; + + const int ncols = dst->ne[0]; + + for (int ir = first; ir < last; ++ir) { + float * y = (float *)dst->data + ir*ncols; + + float sum1 = 0, sum2 = 0; + if (src0->type == GGML_TYPE_F32) { + sum1 = sum_row_squared(ncols, (const float *)src0->data + ir*ncols); + sum2 = sum_row_squared(ncols, (const float *)src2->data + ir*ncols); + } + else if (src0->type == GGML_TYPE_F16) { + sum1 = sum_row_squared(ncols, (const ggml_half *)src0->data + ir*ncols); + sum2 = sum_row_squared(ncols, (const ggml_half *)src2->data + ir*ncols); + } + else { + sum1 = sum_row_squared(ncols, (const ggml_bf16_t *)src0->data + ir*ncols); + sum2 = sum_row_squared(ncols, (const ggml_bf16_t *)src2->data + ir*ncols); + } + + const float mean1 = sum1/ncols; + const float mean2 = sum2/ncols; + const float scale1 = 1.0f/sqrtf(mean1 + eps); + const float scale2 = 1.0f/sqrtf(mean2 + eps); + if (src0->type == GGML_TYPE_F32) { + rms_rms_add(ncols, scale1, scale2, (const float *)src0->data + ir*ncols, (const float *)src2->data + ir*ncols, c1, c2, y); + } + else if (src0->type == GGML_TYPE_F16) { + rms_rms_add(ncols, scale1, scale2, (const ggml_half *)src0->data + ir*ncols, (const ggml_half *)src2->data + ir*ncols, c1, c2, y); + } + else { + rms_rms_add(ncols, scale1, scale2, (const ggml_bf16_t *)src0->data + ir*ncols, (const ggml_bf16_t *)src2->data + ir*ncols, c1, c2, y); + } + } +} diff --git a/ggml/src/iqk/iqk_cpu_ops.h b/ggml/src/iqk/iqk_cpu_ops.h index cd3c641e..0a049caf 100644 --- a/ggml/src/iqk/iqk_cpu_ops.h +++ b/ggml/src/iqk/iqk_cpu_ops.h @@ -39,6 +39,8 @@ bool iqk_ssm_conv4(int nr, int nc, int nt, const float * x0, const float * s0, const float * c, float * dst, float * dst_silu, int ith, int nth); +void iqk_rms_rms_add(struct ggml_tensor * dst, int ith, int nth); + #ifdef __cplusplus } #endif diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index c973697e..82239652 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -6100,18 +6100,15 @@ static ggml_cgraph * build_gemma4_graph_parallel(llm_build_context & llm, llama_ if (il == 0) { sa_inp[id] = inpL; - //sa_inp[id] = do_split_norm(ctx0, inpL, l.attn_norm, hparams, cb, id, il_cb, false); } else { GGML_ASSERT(inpL->op == GGML_OP_REDUCE); auto cur = get_input_tensor_sm_graph(ctx0, inpL, id); if (is_moe) { GGML_ASSERT(inpL_moe && inpL_moe->op == GGML_OP_REDUCE); - cur = do_split_norm(ctx0, cur, model.layers[il-1].ffn_post_norm_1, hparams, cb, id, il_cb, false); - cb(cur, "ffn_post_norm", il_cb); auto cur_moe = get_input_tensor_sm_graph(ctx0, inpL_moe, id); - cur_moe = do_split_norm(ctx0, cur_moe, model.layers[il-1].ffn_post_norm_2, hparams, cb, id, il_cb, false); - cb(cur, "ffn_moe_post_norm", il_cb); - cur = ggml_add(ctx0, cur, cur_moe); + auto post_norm_1 = (ggml_split_tensor_t *)model.layers[il-1].ffn_post_norm_1->extra; + auto post_norm_2 = (ggml_split_tensor_t *)model.layers[il-1].ffn_post_norm_2->extra; + cur = ggml_fused_rms_rms_add(ctx0, cur, post_norm_1->splits[id], cur_moe, post_norm_2->splits[id], hparams.f_norm_rms_eps); cb(cur, "ffn_combined", il_cb); } cur = do_split_norm(ctx0, cur, model.layers[il-1].ffn_post_norm, hparams, cb, id, il_cb, false); @@ -6361,13 +6358,9 @@ static ggml_cgraph * build_gemma4_graph_parallel(llm_build_context & llm, llama_ } auto post_norm_1 = (const ggml_split_tensor_t *)model.layers[hparams.n_layer-1].ffn_post_norm_1->extra; auto post_norm_2 = (const ggml_split_tensor_t *)model.layers[hparams.n_layer-1].ffn_post_norm_2->extra; - cur = llm.llm_build_norm(ctx0, cur, hparams, post_norm_1->splits[idx], NULL, LLM_NORM_RMS, cb, -1); + cur = ggml_fused_rms_rms_add(ctx0, cur, post_norm_1->splits[idx], cur_moe, post_norm_2->splits[idx], hparams.f_norm_rms_eps); cur->op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t) - 1] = 0xff; ggml_build_forward_expand(gf, cur); - cur_moe = llm.llm_build_norm(ctx0, cur_moe, hparams, post_norm_2->splits[idx], NULL, LLM_NORM_RMS, cb, -1); - cb(cur, "ffn_post", hparams.n_layer-1); - cb(cur_moe, "ffn_post_moe", hparams.n_layer-1); - cur = ggml_add(ctx0, cur, cur_moe); cb(cur, "ffn_combined", hparams.n_layer-1); } cur = llm.llm_build_norm(ctx0, cur, hparams, post_norm->splits[idx], NULL, LLM_NORM_RMS, cb, -1); @@ -6568,9 +6561,6 @@ ggml_cgraph * llm_build_context::build_gemma4() { nullptr, LLM_FFN_GELU, LLM_FFN_PAR, cb, il, gf); - cur_mlp = llm_build_norm(ctx0, cur_mlp, hparams, model.layers[il].ffn_post_norm_1, nullptr, LLM_NORM_RMS, cb, il); - cb(cur_mlp, "ffn_mlp", il); - // Expert FFN auto cur_moe = llm_build_norm(ctx0, attn_out, hparams, model.layers[il].ffn_pre_norm_2, nullptr, LLM_NORM_RMS, cb, il); cb(cur_moe, "ffn_norm_2", il); @@ -6596,11 +6586,9 @@ ggml_cgraph * llm_build_context::build_gemma4() { model.layers[il].ffn_up_gate_exps, nullptr, logits, model.layers[il].ffn_down_exps_s); - cur_moe = llm_build_norm(ctx0, cur_moe, hparams, model.layers[il].ffn_post_norm_2, nullptr, LLM_NORM_RMS, cb, il); - cb(cur_moe, "ffn_moe", il); - - cur = ggml_add(ctx0, cur_mlp, cur_moe); + cur = ggml_fused_rms_rms_add(ctx0, cur_mlp, model.layers[il].ffn_post_norm_1, cur_moe, model.layers[il].ffn_post_norm_2, hparams.f_norm_rms_eps); cb(cur, "ffn_moe_combined", il); + ggml_build_forward_expand(gf, cur); cur = llm_build_norm(ctx0, cur, hparams, model.layers[il].ffn_post_norm, NULL, LLM_NORM_RMS, cb, -1); cb(cur, "ffn_post_norm", -1);