Fuse add+add+fused_rms

ggml/src/ggml-cuda.cu

@@ -3129,11 +3129,26 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
             ggml_cuda_dup(ctx, dst);
             break;
         case GGML_OP_ADD:
-            if (i + 1 < cgraph->n_nodes &&
+            if (i + 2 < cgraph->n_nodes &&
+                cgraph->nodes[i+1]->op == GGML_OP_ADD &&
+                cgraph->nodes[i+2]->op == GGML_OP_FUSED_RMS_NORM &&
+                ggml_is_contiguous(dst->src[0]) &&
+                ggml_is_contiguous(dst->src[1]) &&
+                ggml_are_same_shape(dst->src[0], dst->src[1]) &&
+                dst == cgraph->nodes[i+1]->src[0] &&
+                ggml_is_contiguous(cgraph->nodes[i+1]->src[1]) &&
+                ggml_are_same_shape(dst, cgraph->nodes[i+1]->src[1]) &&
+                cgraph->nodes[i+1] == cgraph->nodes[i+2]->src[0]) {
+                //printf("Fusing add->add->fused_rms of %s, %s, %s\n", dst->name, cgraph->nodes[i+1]->name, cgraph->nodes[i+2]->name);
+                ggml_cuda_op_fused_add_add_rms_norm(ctx, dst, cgraph->nodes[i+1], cgraph->nodes[i+2]);
+                i += 2;
+            }
+            else if (i + 1 < cgraph->n_nodes &&
                 cgraph->nodes[i+1]->op == GGML_OP_FUSED_RMS_NORM &&
                 ggml_is_contiguous(dst->src[0]) &&
                 ggml_is_contiguous(dst->src[1]) &&
-                ggml_are_same_shape(dst->src[0], dst->src[1])) {
+                ggml_are_same_shape(dst->src[0], dst->src[1]) &&
+                dst == cgraph->nodes[i+1]->src[0]) {
                 ggml_cuda_op_fused_add_rms_norm(ctx, dst, cgraph->nodes[i+1]);
                 ++i;
             } else {

ggml/src/ggml-cuda/norm.cu

@@ -492,6 +492,41 @@ static __global__ void fused_add_rms_norm_f32(const float * a, const float * b,
     }
 }
 
+template <int block_size>
+static __global__ void fused_add_add_rms_norm_f32(const float * a1, const float * a2, const float * b, const float * c,
+        float * dst_add, float * dst, const int ncols, const float eps) {
+    const int row = blockIdx.x*blockDim.y + threadIdx.y;
+    const int tid = threadIdx.x;
+
+    float tmp = 0.0f; // partial sum for thread in warp
+
+    for (int col = tid; col < ncols; col += block_size) {
+        const float xi = a1[row*ncols + col] + a2[row*ncols + col] + b[row*ncols + col];
+        tmp += xi * xi;
+        dst_add[row*ncols + col] = xi;
+    }
+
+    // sum up partial sums
+    tmp = warp_reduce_sum(tmp);
+    if (block_size > WARP_SIZE) {
+        __shared__ float s_sum[32];
+        int warp_id = threadIdx.x / WARP_SIZE;
+        int lane_id = threadIdx.x % WARP_SIZE;
+        if (lane_id == 0) {
+            s_sum[warp_id] = tmp;
+        }
+        __syncthreads();
+        tmp = lane_id < block_size/WARP_SIZE ? s_sum[lane_id] : 0.0f;
+        tmp = warp_reduce_sum(tmp);
+    }
+
+    const float mean = tmp / ncols;
+    const float scale = rsqrtf(mean + eps);
+
+    for (int col = tid; col < ncols; col += block_size) {
+        dst[row*ncols + col] = scale * c[col] * dst_add[row*ncols + col];
+    }
+}
 
 static void fused_add_rms_norm_f32_cuda(const float * a, const float * b, const float * c, float * dst_add, float * dst,
         const int ncols, const int nrows, const float eps, cudaStream_t stream) {
@@ -538,3 +573,49 @@ void ggml_cuda_op_fused_add_rms_norm(ggml_backend_cuda_context & ctx, ggml_tenso
             src1_d, (float *)add->data, dst_d, ne00, nrows, eps, stream);
 }
 
+static void fused_add_add_rms_norm_f32_cuda(const float * a1, const float * a2, const float * b, const float * c, float * dst_add, float * dst,
+        const int ncols, const int nrows, const float eps, cudaStream_t stream) {
+    GGML_ASSERT(ncols % WARP_SIZE == 0);
+    if (ncols < 1024) {
+        const dim3 block_dims(256, 1, 1);
+        fused_add_add_rms_norm_f32<256><<<nrows, block_dims, 0, stream>>>(a1, a2, b, c, dst_add, dst, ncols, eps);
+    } else {
+        const dim3 block_dims(1024, 1, 1);
+        fused_add_add_rms_norm_f32<1024><<<nrows, block_dims, 0, stream>>>(a1, a2, b, c, dst_add, dst, ncols, eps);
+    }
+}
+
+void ggml_cuda_op_fused_add_add_rms_norm(ggml_backend_cuda_context & ctx,
+        ggml_tensor * add1, ggml_tensor * add2, ggml_tensor * dst) {
+
+    const ggml_tensor * src0 = dst->src[0];
+    const ggml_tensor * src1 = dst->src[1];
+    //const float * src0_d = (const float *)src0->data;
+    const float * src1_d = (const float *)src1->data;
+    float * dst_d = (float *)dst->data;
+    cudaStream_t stream = ctx.stream();
+
+    GGML_ASSERT(add1->data == add2->src[0]->data);
+    GGML_ASSERT(add2->data == src0->data);
+    GGML_ASSERT(ggml_is_contiguous(src0));
+    //GGML_ASSERT(ggml_is_contiguous(add->src[0]));
+    //GGML_ASSERT(ggml_is_contiguous(add->src[1]));
+    //GGML_ASSERT(ggml_are_same_shape(add->src[0], add->src[1]));
+    //GGML_ASSERT(ggml_are_same_shape(add->src[0], src0));
+    //GGML_ASSERT(add->src[0]->type == GGML_TYPE_F32);
+    //GGML_ASSERT(add->src[1]->type == GGML_TYPE_F32);
+    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT(src1->type == GGML_TYPE_F32);
+    GGML_ASSERT( dst->type == GGML_TYPE_F32);
+    GGML_ASSERT(src0->ne[0] == src1->ne[0]);
+    GGML_ASSERT(ggml_nrows(src1) == 1);
+
+    float eps;
+    memcpy(&eps, dst->op_params, sizeof(float));
+
+    const int64_t ne00 = src0->ne[0];
+
+    const int64_t nrows = ggml_nrows(src0);
+    fused_add_add_rms_norm_f32_cuda((const float *)add1->src[0]->data, (const float *)add1->src[1]->data, (const float *)add2->src[1]->data,
+            src1_d, (float *)add2->data, dst_d, ne00, nrows, eps, stream);
+}

ggml/src/ggml-cuda/norm.cuh

@@ -9,3 +9,5 @@ void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
 void ggml_cuda_op_fused_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
 
 void ggml_cuda_op_fused_add_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * add, ggml_tensor * dst);
+
+void ggml_cuda_op_fused_add_add_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * add1, ggml_tensor * add2, ggml_tensor * dst);

src/llama-build-context.cpp

@@ -7793,6 +7793,7 @@ ggml_cgraph * llm_build_context::build_openai_moe() {
 
         cur = ffn_inp;
         cur = llm_build_norm(ctx0, cur,  hparams, model.layers[il].attn_post_norm, nullptr, LLM_NORM_RMS, cb, il);
+        ggml_build_forward_expand(gf, cur);
         cb(cur, "attn_post_norm", il);
 
         bool use_dup_bias = cur->ne[1] < 32 && model.layers[il].ffn_up_exps_b_dup &&