refactor: replace manual gated activations with native ggml GLU (#2250)

src/core/ggml_extend.cpp

@@ -135,21 +135,11 @@ ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_firs
     // x: [ne3, ne2, ne1, ne0]
     // return: [ne3, ne2, ne1, ne0/2]
 
-    auto x_vec = ggml_ext_chunk(ctx, x, 2, 0, false);
-    ggml_tensor* gate;
-    if (gate_first) {
-        gate = x_vec[0];
-        x    = x_vec[1];
-    } else {
-        x    = x_vec[0];
-        gate = x_vec[1];
+    GGML_ASSERT(x->ne[0] % 2 == 0);
+    if (!ggml_is_contiguous_1(x)) {
+        x = ggml_cont(ctx, x);
     }
-    gate = ggml_cont(ctx, gate);
-    gate = ggml_silu_inplace(ctx, gate);
-
-    x = ggml_mul(ctx, x, gate);  // [ne3, ne2, ne1, ne0/2]
-
-    return x;
+    return gate_first ? ggml_swiglu(ctx, x) : ggml_swiglu_swapped(ctx, x);
 }
 
 ggml_tensor* ggml_ext_group_norm_32(ggml_context* ctx,

src/model/diffusion/ernie_image.hpp

@@ -204,9 +204,8 @@ namespace ErnieImage {
             auto linear_fc2 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc2"]);
 
             auto gate = gate_proj->forward(ctx, x);
-            gate      = ggml_ext_gelu(ctx->ggml_ctx, gate);
             x         = up_proj->forward(ctx, x);
-            x         = ggml_mul(ctx->ggml_ctx, x, gate);
+            x         = ggml_geglu_split(ctx->ggml_ctx, gate, x);
             x         = linear_fc2->forward(ctx, x);
             return x;
         }

src/model/diffusion/flux.hpp

@@ -358,9 +358,8 @@ namespace Flux {
             auto down_proj = std::dynamic_pointer_cast<Linear>(blocks["down_proj"]);
 
             auto gate = gate_proj->forward(ctx, x);
-            gate      = ggml_silu_inplace(ctx->ggml_ctx, gate);
             x         = up_proj->forward(ctx, x);
-            x         = ggml_mul(ctx->ggml_ctx, x, gate);
+            x         = ggml_swiglu_split(ctx->ggml_ctx, gate, x);
             x         = down_proj->forward(ctx, x);
             return x;
         }
@@ -850,12 +849,10 @@ namespace Flux {
             auto res_x = x;
             x          = norm->forward(ctx, x);  // [batch_size, n_token, hidden_size_x]
 
-            auto x1 = ggml_mul_mat(ctx->ggml_ctx, fc1_gate, x);  // [batch_size, n_token, hidden_size_x*mlp_ratio]
-            x1      = ggml_silu_inplace(ctx->ggml_ctx, x1);
-
+            auto x1 = ggml_mul_mat(ctx->ggml_ctx, fc1_gate, x);   // [batch_size, n_token, hidden_size_x*mlp_ratio]
             auto x2 = ggml_mul_mat(ctx->ggml_ctx, fc1_value, x);  // [batch_size, n_token, hidden_size_x*mlp_ratio]
 
-            x = ggml_mul_inplace(ctx->ggml_ctx, x1, x2);  // [batch_size, n_token, hidden_size_x*mlp_ratio]
+            x = ggml_swiglu_split(ctx->ggml_ctx, x1, x2);  // [batch_size, n_token, hidden_size_x*mlp_ratio]
 
             x = ggml_mul_mat(ctx->ggml_ctx, fc2, x);  // [batch_size, n_token, hidden_size_x]
 

src/model/diffusion/ideogram4.hpp

@@ -244,9 +244,9 @@ namespace Ideogram4 {
             auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
             auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
 
-            auto x1 = ggml_silu(ctx->ggml_ctx, w1->forward(ctx, x));
+            auto x1 = w1->forward(ctx, x);
             auto x3 = w3->forward(ctx, x);
-            x       = ggml_mul(ctx->ggml_ctx, x1, x3);
+            x       = ggml_swiglu_split(ctx->ggml_ctx, x1, x3);
             x       = w2->forward(ctx, x);
             return x;
         }

src/model/diffusion/krea2.hpp

@@ -208,9 +208,9 @@ namespace Krea2 {
             auto up   = std::dynamic_pointer_cast<Linear>(blocks["up"]);
             auto down = std::dynamic_pointer_cast<Linear>(blocks["down"]);
 
-            auto gated = ggml_silu(ctx->ggml_ctx, gate->forward(ctx, x));
-            auto up_x  = up->forward(ctx, x);
-            x          = ggml_mul(ctx->ggml_ctx, gated, up_x);
+            auto gate_x = gate->forward(ctx, x);
+            auto up_x   = up->forward(ctx, x);
+            x           = ggml_swiglu_split(ctx->ggml_ctx, gate_x, up_x);
             return down->forward(ctx, x);
         }
     };

src/model/diffusion/lens.hpp

@@ -104,9 +104,9 @@ namespace Lens {
             auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
             auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
 
-            auto gate = ggml_silu(ctx->ggml_ctx, w1->forward(ctx, x));
+            auto gate = w1->forward(ctx, x);
             auto up   = w3->forward(ctx, x);
-            x         = ggml_mul(ctx->ggml_ctx, gate, up);
+            x         = ggml_swiglu_split(ctx->ggml_ctx, gate, up);
             return w2->forward(ctx, x);
         }
     };

src/model/diffusion/lingbot_video.hpp

@@ -231,9 +231,8 @@ namespace LingBotVideo {
             auto down_proj = std::dynamic_pointer_cast<Linear>(blocks["down_proj"]);
 
             auto gate = gate_proj->forward(ctx, x);
-            gate      = ggml_silu_inplace(ctx->ggml_ctx, gate);
             auto up   = up_proj->forward(ctx, x);
-            x         = ggml_mul(ctx->ggml_ctx, gate, up);
+            x         = ggml_swiglu_split(ctx->ggml_ctx, gate, up);
             x         = down_proj->forward(ctx, x);
             return x;
         }
@@ -406,9 +405,8 @@ namespace LingBotVideo {
 
             x         = ggml_reshape_3d(ctx->ggml_ctx, x, hidden_size, 1, n_token_total);
             auto gate = expert_linear(ctx, "experts.w1", x, selected_experts);
-            gate      = ggml_silu_inplace(ctx->ggml_ctx, gate);
             auto up   = expert_linear(ctx, "experts.w3", x, selected_experts);
-            auto act  = ggml_mul(ctx->ggml_ctx, gate, up);
+            auto act  = ggml_swiglu_split(ctx->ggml_ctx, gate, up);
             auto out  = expert_linear(ctx, "experts.w2", act, selected_experts);
             out       = ggml_mul(ctx->ggml_ctx, out, weights);
 

src/model/diffusion/minit2i.hpp

@@ -207,9 +207,9 @@ namespace MiniT2I {
             auto w1   = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
             auto w3   = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
             auto w2   = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
-            auto gate = ggml_silu(ctx->ggml_ctx, w1->forward(ctx, x));
+            auto gate = w1->forward(ctx, x);
             auto up   = w3->forward(ctx, x);
-            return w2->forward(ctx, ggml_mul(ctx->ggml_ctx, gate, up));
+            return w2->forward(ctx, ggml_swiglu_split(ctx->ggml_ctx, gate, up));
         }
     };
 

src/model/diffusion/pid.hpp

@@ -253,8 +253,8 @@ namespace Pid {
             auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
             auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
             auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
-            auto h  = ggml_silu_inplace(ctx->ggml_ctx, w1->forward(ctx, x));
-            h       = ggml_mul_inplace(ctx->ggml_ctx, h, w3->forward(ctx, x));
+            auto h  = w1->forward(ctx, x);
+            h       = ggml_swiglu_split(ctx->ggml_ctx, h, w3->forward(ctx, x));
             return w2->forward(ctx, h);
         }
     };

src/model/te/llm.hpp

@@ -517,13 +517,13 @@ namespace LLM {
             auto up_proj   = std::dynamic_pointer_cast<Linear>(blocks["up_proj"]);
             auto down_proj = std::dynamic_pointer_cast<Linear>(blocks["down_proj"]);
 
-            auto h = gate_proj->forward(ctx, x);
+            auto h  = gate_proj->forward(ctx, x);
+            auto up = up_proj->forward(ctx, x);
             if (activation == MLPActivation::GELU_TANH) {
-                h = ggml_ext_gelu(ctx->ggml_ctx, h, true);
+                h = ggml_geglu_split(ctx->ggml_ctx, h, up);
             } else {
-                h = ggml_silu_inplace(ctx->ggml_ctx, h);
+                h = ggml_swiglu_split(ctx->ggml_ctx, h, up);
             }
-            h = ggml_mul_inplace(ctx->ggml_ctx, h, up_proj->forward(ctx, x));
             h = down_proj->forward(ctx, h);
             return h;
         }
@@ -1027,8 +1027,7 @@ namespace LLM {
                 auto up_proj   = std::dynamic_pointer_cast<Linear>(blocks["up_proj"]);
                 auto down_proj = std::dynamic_pointer_cast<Linear>(blocks["down_proj"]);
                 auto h         = gate_proj->forward(ctx, x);
-                h              = ggml_silu_inplace(ctx->ggml_ctx, h);
-                h              = ggml_mul_inplace(ctx->ggml_ctx, h, up_proj->forward(ctx, x));
+                h              = ggml_swiglu_split(ctx->ggml_ctx, h, up_proj->forward(ctx, x));
                 x              = down_proj->forward(ctx, h);
             }
             return x;
@@ -1520,7 +1519,7 @@ namespace LLM {
             out        = ggml_view_3d(c, out, head_dim, value_heads, n, head_dim * sizeof(float), head_dim * value_heads * sizeof(float), 0);
             out        = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["norm"])->forward(ctx, out);
             auto z     = ggml_reshape_3d(c, linear("in_proj_z"), head_dim, value_heads, n);
-            out        = ggml_mul(c, out, ggml_silu(c, z));
+            out        = ggml_swiglu_split(c, z, out);
             out        = ggml_reshape_3d(c, out, head_dim * value_heads, n, 1);
             return std::dynamic_pointer_cast<Linear>(blocks["out_proj"])->forward(ctx, out);
         }

src/model/te/t5.hpp

@@ -156,9 +156,9 @@ public:
         auto wi_1 = std::dynamic_pointer_cast<Linear>(blocks["wi_1"]);
         auto wo   = std::dynamic_pointer_cast<Linear>(blocks["wo"]);
 
-        auto hidden_gelu   = ggml_ext_gelu(ctx->ggml_ctx, wi_0->forward(ctx, x), true);
+        auto hidden_gate   = wi_0->forward(ctx, x);
         auto hidden_linear = wi_1->forward(ctx, x);
-        x                  = ggml_mul_inplace(ctx->ggml_ctx, hidden_gelu, hidden_linear);
+        x                  = ggml_geglu_split(ctx->ggml_ctx, hidden_gate, hidden_linear);
         x                  = wo->forward(ctx, x);
         return x;
     }

src/model/vae/minimax_h3_audio_vae.hpp

@@ -134,8 +134,8 @@ namespace MiniMaxH3 {
             auto w1   = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
             auto w2   = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
             x         = norm->forward(ctx, x);
-            auto gate = ggml_ext_gelu(ctx->ggml_ctx, w0->forward(ctx, x), true);
-            return w2->forward(ctx, ggml_mul(ctx->ggml_ctx, gate, w1->forward(ctx, x)));
+            auto gate = w0->forward(ctx, x);
+            return w2->forward(ctx, ggml_geglu_split(ctx->ggml_ctx, gate, w1->forward(ctx, x)));
         }
     };
 

src/model/vae/minimax_h3_vae.hpp

@@ -316,11 +316,8 @@ namespace MiniMaxH3VAE {
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
             auto w1   = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
             auto w2   = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
-            auto gate = ggml_ext_chunk(ctx->ggml_ctx, w1->forward(ctx, x), 2, 0);
-            return w2->forward(ctx,
-                               ggml_mul(ctx->ggml_ctx,
-                                        ggml_silu(ctx->ggml_ctx, gate[0]),
-                                        gate[1]));
+            auto gate = ggml_swiglu(ctx->ggml_ctx, w1->forward(ctx, x));
+            return w2->forward(ctx, gate);
         }
     };