Fuse SILU and SSM_CONV (CPU) (#1421)

ggml/src/ggml.c

@@ -22232,9 +22232,9 @@ static void ggml_compute_forward_flash_attn_back(
 
 // ggml_compute_forward_ssm_conv
 
-static void ggml_compute_forward_ssm_conv_f32(
+static int ggml_compute_forward_ssm_conv_f32(
         const struct ggml_compute_params * params,
-        struct ggml_tensor * dst) {
+        struct ggml_tensor * dst, int node, const struct ggml_cgraph * cgraph) {
     const struct ggml_tensor * src0 = dst->src[0]; // conv_state
     const struct ggml_tensor * src1 = dst->src[1]; // x
     const struct ggml_tensor * src2 = dst->src[2]; // conv1d.weight
@@ -22258,10 +22258,17 @@ static void ggml_compute_forward_ssm_conv_f32(
     GGML_ASSERT(src2->nb[2] == src2->ne[1]*src2->ne[0]*sizeof(float));
 
     if (n_kv == 1 && nc == 4) {
+        float * dst_silu = NULL;
+        if (node < cgraph->n_nodes + 2 &&
+            cgraph->nodes[node+1]->op == GGML_OP_VIEW && cgraph->nodes[node+1]->src[0] == dst &&
+            cgraph->nodes[node+2]->op == GGML_OP_UNARY && cgraph->nodes[node+2]->src[0] == cgraph->nodes[node+1] &&
+            (enum ggml_unary_op)cgraph->nodes[node+2]->op_params[0] == GGML_UNARY_OP_SILU) {
+            dst_silu = (float *)cgraph->nodes[node+2]->data;
+        }
         if (iqk_ssm_conv4(nr, nc, n_t, src0->nb[1], src1->nb[0], src1->nb[1], src2->nb[1],
                     (const float *)src1->data, (const float *)src0->data, (const float *)src2->data,
-                    (float *)dst->data, ith, nth)) {
-            return;
+                    (float *)dst->data, dst_silu, ith, nth)) {
+            return node + (dst_silu ? 2 : 0);
         }
     }
 
@@ -22343,15 +22350,15 @@ static void ggml_compute_forward_ssm_conv_f32(
             x[i1] = sumf;
         }
     }
+    return node;
 }
 
-static void ggml_compute_forward_ssm_conv(
-        const struct ggml_compute_params * params,
-        struct ggml_tensor * dst) {
+static int ggml_compute_forward_ssm_conv(const struct ggml_compute_params * params,
+        struct ggml_tensor * dst, int node, const struct ggml_cgraph * cgraph) {
     switch (dst->src[0]->type) {
         case GGML_TYPE_F32:
             {
-                ggml_compute_forward_ssm_conv_f32(params, dst);
+                return ggml_compute_forward_ssm_conv_f32(params, dst, node, cgraph);
             } break;
         default:
             {
@@ -24407,7 +24414,7 @@ static int ggml_compute_forward(struct ggml_compute_params * params, struct ggml
             } break;
         case GGML_OP_SSM_CONV:
             {
-                ggml_compute_forward_ssm_conv(params, tensor);
+                i = ggml_compute_forward_ssm_conv(params, tensor, i, cgraph);
             } break;
         case GGML_OP_SSM_SCAN:
             {

ggml/src/iqk/iqk_cpu_ops.cpp

@@ -554,7 +554,7 @@ float iqk_exp_with_thresh(int n, float * logits, float max, float min) {
 bool iqk_ssm_conv4(int nr, int nc, int nt,
         uint64_t nb01, uint64_t nb10, uint64_t nb11, uint64_t nb21,
         const float * x0_in, const float * s0_in, const float * c_in,
-        float * dst, int ith, int nth) {
+        float * dst, float * dst_silu, int ith, int nth) {
 #ifdef __AVX2__
     if (nt <= 32 || nc != 4 || nr%16 != 0) {
         return false;
@@ -566,7 +566,7 @@ bool iqk_ssm_conv4(int nr, int nc, int nt,
     __m256 vs[8], vc[8];
     float aux[64];
     for (int ir = ir0; ir < ir1; ++ir) {
-        auto x  = dst   + 16*ir;
+        auto x  = dst_silu == nullptr ? dst + 16*ir : dst_silu + 16*ir;
         auto s  = dst   + 16*ir*nb21/sizeof(float) + nr*nt;
         auto s0 = s0_in + 16*ir*nb01/sizeof(float); // {d_conv - 1, d_inner, n_kv}
         auto x0 = x0_in + 16*ir*nb10/sizeof(float);
@@ -601,6 +601,10 @@ bool iqk_ssm_conv4(int nr, int nc, int nt,
                 sum1 = _mm256_fmadd_ps(vs[ii+0], vc[k+0], sum1);
                 sum2 = _mm256_fmadd_ps(vs[ii+4], vc[k+4], sum2);
             }
+            if (dst_silu) {
+                sum1 = v_silu(sum1);
+                sum2 = v_silu(sum2);
+            }
             _mm256_storeu_ps(x+0, sum1);
             _mm256_storeu_ps(x+8, sum2);
             x0 += nb11/sizeof(float);

ggml/src/iqk/iqk_cpu_ops.h

@@ -35,7 +35,7 @@ float iqk_exp_with_thresh(int n, float * logits, float max, float min);
 bool iqk_ssm_conv4(int nr, int nc, int nt,
         uint64_t nb01, uint64_t nb10, uint64_t nb11, uint64_t nb21,
         const float * x0, const float * s0, const float * c,
-        float * dst, int ith, int nth);
+        float * dst, float * dst_silu, int ith, int nth);
 
 #ifdef __cplusplus
 }

src/llama-delta-net.cpp

@@ -341,6 +341,7 @@ ggml_tensor * delta_net::build_qkv(ggml_context * ctx0, ggml_tensor * state_stor
     ggml_tensor * conv_output = ggml_view_2d(ctx0, conv_output_raw, conv_dim, n_tok, conv_dim * ggml_element_size(conv_output_raw), 0);
     ggml_tensor * conv_output_silu = ggml_silu(ctx0, conv_output);
     cb(conv_output_silu, "conv_output_silu", il);
+    ggml_build_forward_expand(gf, conv_output_silu);
 
     // Calculate the total conv dimension
     int64_t qkv_dim = head_k_dim * num_k_heads * 2 + head_v_dim * num_v_heads;