ggml-cpu: fuse RMS_NORM + MUL on CPU backend

Add a fused rms_norm + mul kernel that computes the output in a single pass,
avoiding materialization of the intermediate rms_norm result. The graph compute
loop detects the pattern and dispatches to the fused path. fusion can be
disabled via the GGML_CPU_DISABLE_FUSION environment variable.

Author: zzzzwc

ggml/src/ggml-cpu/ggml-cpu.c

@@ -2959,6 +2959,53 @@ struct ggml_cplan ggml_graph_plan(
     return cplan;
 }
 
+
+// Read env GGML_CPU_DISABLE_FUSION once and cache it
+static bool get_disable_fusion(void) {
+    static bool value = false;
+    static bool initialized = false;
+    if (!initialized) {
+        value = (getenv("GGML_CPU_DISABLE_FUSION") != NULL);
+        initialized = true;
+    }
+    return value;
+}
+
+// Try to fuse the current node with subsequent nodes for better performance.
+// Returns true if fusion was applied.
+static bool ggml_cpu_try_fuse_ops(
+        const struct ggml_cgraph * cgraph,
+        const int node_n,
+        const struct ggml_compute_params * params,
+        const struct ggml_cplan * cplan) {
+    if (get_disable_fusion() || cplan->use_ref) {
+        return false;
+    }
+
+    struct ggml_tensor * node = cgraph->nodes[node_n];
+
+    if (node->op == GGML_OP_RMS_NORM) {
+        // RMS_NORM + MUL fusion
+        const enum ggml_op fuse_ops[] = { GGML_OP_RMS_NORM, GGML_OP_MUL };
+        if (ggml_can_fuse(cgraph, node_n, fuse_ops, 2)) {
+            struct ggml_tensor * mul_node = cgraph->nodes[node_n + 1];
+            const struct ggml_tensor * mul_w = (mul_node->src[0] == node)
+                ? mul_node->src[1] : mul_node->src[0];
+            if (node->src[0]->type  == GGML_TYPE_F32 &&
+                mul_node->type      == GGML_TYPE_F32 &&
+                mul_w->type         == GGML_TYPE_F32 &&
+                mul_w->ne[0]        == node->ne[0]   &&
+                mul_w->nb[0]        == sizeof(float)) {
+
+                ggml_compute_forward_rms_norm_mul_fused(params, node, mul_node);
+                return true;
+            }
+        }
+    }
+
+    return false;
+}
+
 static thread_ret_t ggml_graph_compute_thread(void * data) {
     struct ggml_compute_state * state = (struct ggml_compute_state *) data;
     struct ggml_threadpool    * tp    = state->threadpool;
@@ -2995,7 +3042,12 @@ static thread_ret_t ggml_graph_compute_thread(void * data) {
             continue;
         }
 
-        ggml_compute_forward(&params, node);
+        // Try fused ops, fall back to normal compute
+        if (ggml_cpu_try_fuse_ops(cgraph, node_n, &params, cplan)) {
+            node_n++;
+        } else {
+            ggml_compute_forward(&params, node);
+        }
 
         if (state->ith == 0 && cplan->abort_callback &&
                 cplan->abort_callback(cplan->abort_callback_data)) {

ggml/src/ggml-cpu/ops.cpp

@@ -3782,6 +3782,68 @@ void ggml_compute_forward_rms_norm(
     }
 }
 
+// Fused RMS_NORM + MUL: computes dst = rms_norm(src0) * src1 in a single pass.
+// This avoids materializing the intermediate rms_norm result in memory.
+void ggml_compute_forward_rms_norm_mul_fused(
+        const ggml_compute_params * params,
+        ggml_tensor * rms_norm_dst,
+        ggml_tensor * mul_dst) {
+
+    //   src0 = rms_norm input, src1 = mul weight, dst = mul output
+    const ggml_tensor * src0 = rms_norm_dst->src[0];
+    const ggml_tensor * src1 = (mul_dst->src[0] == rms_norm_dst) ? mul_dst->src[1] : mul_dst->src[0];
+    ggml_tensor       * dst  = mul_dst;
+    GGML_ASSERT(mul_dst->src[0] == rms_norm_dst || mul_dst->src[1] == rms_norm_dst);
+
+    GGML_ASSERT(ggml_are_same_shape(src0, dst));
+    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT(src1->type == GGML_TYPE_F32);
+    GGML_ASSERT( dst->type == GGML_TYPE_F32);
+    GGML_ASSERT(src0->nb[0] == sizeof(float));
+    GGML_ASSERT(src1->nb[0] == sizeof(float));
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    GGML_TENSOR_BINARY_OP_LOCALS
+
+    GGML_ASSERT(ne10 == ne00); // no column broadcasting on the weight
+
+    float eps;
+    memcpy(&eps, rms_norm_dst->op_params, sizeof(float));
+    GGML_ASSERT(eps >= 0.0f);
+    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);
+
+                ggml_float sum = 0.0;
+                // worth switching to explicit SIMD?
+                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);
+
+                assert(scale > 0.0f);
+
+                const int64_t i11 = i01 % ne11;
+                const int64_t i12 = i02 % ne12;
+                const int64_t i13 = i03 % ne13;
+                const float * w = (float *) ((char *) src1->data + i11*nb11 + i12*nb12 + i13*nb13);
+
+                float * y = (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3);
+
+                // worth switching to explicit SIMD?
+                for (int64_t i00 = 0; i00 < ne00; i00++) {
+                    y[i00] = x[i00] * scale * w[i00];
+                }
+            }
+        }
+    }
+}
+
 static void ggml_compute_forward_rms_norm_back_f32(
         const ggml_compute_params * params,
         ggml_tensor * dst) {

ggml/src/ggml-cpu/ops.h

@@ -44,6 +44,7 @@ void ggml_compute_forward_concat(const struct ggml_compute_params * params, stru
 void ggml_compute_forward_silu_back(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_rms_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_rms_norm_mul_fused(const struct ggml_compute_params * params, struct ggml_tensor * rms_norm_dst, struct ggml_tensor * mul_dst);
 void ggml_compute_forward_rms_norm_back(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_group_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_l2_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst);