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

Faster fused_rms_norm on the CPU (#1427)

This commit is contained in:
Kawrakow 2026-03-14 16:33:34 +01:00 committed by GitHub
parent be2940f57a
commit aa053205e8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 20 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
ggml/src/ggml.c
View File

@ -16402,29 +16402,31 @@ static void ggml_compute_forward_fused_rms_norm_f32(
GGML_ASSERT(eps > 0.0f);
// TODO: optimize
for (int64_t i03 = 0; i03 < ne03; i03++) {
for (int64_t i02 = 0; i02 < ne02; i02++) {
for (int64_t i01 = ith; i01 < ne01; i01 += nth) {
const float * x = (float *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03);
int nrows = ne03*ne02*ne01;
int nrows_per_thread = (nrows + nth - 1)/nth;
int first = ith*nrows_per_thread;
int last = MIN(nrows, first + nrows_per_thread);
ggml_float sum = 0.0;
for (int64_t i00 = 0; i00 < ne00; i00++) {
sum += (ggml_float)(x[i00] * x[i00]);
}
const float * c = (float *) src1->data;
const float mean = sum/ne00;
for (int ir = first; ir < last; ++ir) {
int i03 = ir/(ne01*ne02);
int i02 = (ir - i03*ne01*ne02)/ne01;
int i01 = ir - i03*ne01*ne02 - i02*ne01;
const float * x = (float *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03);
float * y = (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3);
float * y = (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3);
const float scale = 1.0f/sqrtf(mean + eps);
ggml_vec_mul_f32(ne00, y, x, (const float *)src1->data);
ggml_vec_scale_f32(ne00, y, scale);
}
ggml_float sum = 0.0;
for (int64_t i00 = 0; i00 < ne00; i00++) {
sum += (ggml_float)(x[i00] * x[i00]);
}
const float mean = sum/ne00;
const float scale = 1.0f/sqrtf(mean + eps);
for (int j = 0; j < (int)ne00; ++j) {
y[j] = scale * c[j] * x[j];
}
}
}
static void ggml_compute_forward_fused_rms_norm(