LLM-Inference/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-08-12 22:29:39 +04:00
Code Issues Packages Projects Releases Wiki Activity

Fused fused_rms_norm + fused_rms_norm + add (#1627)

Browse Source 
* Fuse fused_rms + fused_rms + add

* Dedicated fused_rms_norm + fused_rms_norm + add op

* Cleanup
This commit is contained in:
Kawrakow
2026-04-13 13:24:39 +02:00
committed by GitHub
parent 191b53c2cd
commit 55d3c05bf7
8 changed files with 295 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/include/ggml.h
+9
View File
@@ -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(
ggml/src/ggml-cuda.cu
+4
View File
@@ -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:
ggml/src/ggml-cuda/norm.cu
+101
View File
@@ -951,3 +951,104 @@ void ggml_cuda_op_fused_rms_rms_norm([[maybe_unused]] ggml_backend_cuda_context
}
template <int block_size, typename src_t>
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 <typename src_t>
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><<<nrows, block_dims, 0, stream>>>(ncols, nrows, dst, x1, c1, x2, c2, eps);
} else {
const dim3 block_dims(1024, 1, 1);
fused_rms_rms_add_f32<1024><<<nrows, block_dims, 0, stream>>>(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");
}
}
ggml/src/ggml-cuda/norm.cuh
+2
View File
@@ -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);
ggml/src/ggml.c
+41 -2
View File
@@ -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:
ggml/src/iqk/iqk_cpu_ops.cpp
+130 -1
View File
@@ -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);
}
}
}
ggml/src/iqk/iqk_cpu_ops.h
+2
View File
@@ -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
src/llama-build-context.cpp
+6 -18
View File
@@ -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);