From 46018f89edd3e58c60aa8b0207d46448c4c4c7aa Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Sat, 14 Mar 2026 08:27:32 +0100 Subject: [PATCH] Fuse SILU and SSM_CONV (CPU) (#1421) --- ggml/src/ggml.c | 25 ++++++++++++++++--------- ggml/src/iqk/iqk_cpu_ops.cpp | 8 ++++++-- ggml/src/iqk/iqk_cpu_ops.h | 2 +- src/llama-delta-net.cpp | 1 + 4 files changed, 24 insertions(+), 12 deletions(-) diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 16728f74..d11cac99 100644 --- a/ggml/src/ggml.c +++ b/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: { diff --git a/ggml/src/iqk/iqk_cpu_ops.cpp b/ggml/src/iqk/iqk_cpu_ops.cpp index 1cc128b5..0a214f2d 100644 --- a/ggml/src/iqk/iqk_cpu_ops.cpp +++ b/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); diff --git a/ggml/src/iqk/iqk_cpu_ops.h b/ggml/src/iqk/iqk_cpu_ops.h index 6f4317f9..4b98d885 100644 --- a/ggml/src/iqk/iqk_cpu_ops.h +++ b/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 } diff --git a/src/llama-delta-net.cpp b/src/llama-delta-net.cpp index 7fadddeb..50ab7220 100644 --- a/src/llama-delta-net.cpp +++ b/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;