perf: use native SiTU for MOSS with composite fallback (#2210)

src/core/ggml_extend.cpp

@@ -144,6 +144,27 @@ ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_firs
     return gate_first ? ggml_swiglu(ctx, x) : ggml_swiglu_swapped(ctx, x);
 }
 
+ggml_tensor* ggml_ext_situ(GGMLRunnerContext* ctx, ggml_tensor* hidden, ggml_tensor* gate) {
+#ifndef SD_USE_UPSTREAM_GGML
+    if (ctx->backend != nullptr) {
+        auto out = gate ? ggml_glu_split(ctx->ggml_ctx, gate, hidden, GGML_GLU_OP_SITU)
+                        : ggml_glu(ctx->ggml_ctx, hidden, GGML_GLU_OP_SITU, true);
+        if (ggml_backend_supports_op(ctx->backend, out)) {
+            return out;
+        }
+    }
+#endif
+    if (!gate) {
+        const auto dim = hidden->ne[0] / 2;
+        gate           = ggml_ext_slice(ctx->ggml_ctx, hidden, 0, dim, 2 * dim);
+        hidden         = ggml_ext_slice(ctx->ggml_ctx, hidden, 0, 0, dim);
+    }
+    auto a = ggml_scale(ctx->ggml_ctx, ggml_tanh(ctx->ggml_ctx, ggml_scale(ctx->ggml_ctx, gate, 0.25f)), 4.f);
+    a      = ggml_mul(ctx->ggml_ctx, a, ggml_sigmoid(ctx->ggml_ctx, gate));
+    auto b = ggml_scale(ctx->ggml_ctx, ggml_tanh(ctx->ggml_ctx, ggml_scale(ctx->ggml_ctx, hidden, 1.f / 25.f)), 25.f);
+    return ggml_mul(ctx->ggml_ctx, a, b);
+}
+
 ggml_tensor* ggml_ext_group_norm_32(ggml_context* ctx,
                                     ggml_tensor* a) {
     const float eps = 1e-6f;  // default eps parameter

src/core/ggml_extend.h

@@ -52,6 +52,9 @@ std::vector<ggml_tensor*> ggml_ext_chunk(ggml_context* ctx,
 
 ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_first = true);
 
+// Without gate, hidden contains [hidden, gate] packed along ne[0].
+ggml_tensor* ggml_ext_situ(GGMLRunnerContext* ctx, ggml_tensor* hidden, ggml_tensor* gate = nullptr);
+
 ggml_tensor* ggml_ext_group_norm_32(ggml_context* ctx,
                                     ggml_tensor* a);
 

src/model/diffusion/moss_image.hpp

@@ -4,6 +4,7 @@
 #include <algorithm>
 #include <array>
 #include <cinttypes>
+
 #include "core/pe_cache.h"
 
 #include "model/common/ggml_block.hpp"
@@ -55,13 +56,6 @@ namespace MOSSImage {
         }
     };
 
-    inline ggml_tensor* situ(ggml_context* c, ggml_tensor* hidden, ggml_tensor* gate) {
-        auto a = ggml_scale(c, ggml_tanh(c, ggml_scale(c, gate, 0.25f)), 4.f);
-        a      = ggml_mul(c, a, ggml_sigmoid(c, gate));
-        auto b = ggml_scale(c, ggml_tanh(c, ggml_scale(c, hidden, 1.f / 25.f)), 25.f);
-        return ggml_mul(c, a, b);
-    }
-
     struct TextBlock : public GGMLBlock {
         MOSSImageConfig config;
 
@@ -94,7 +88,7 @@ namespace MOSSImage {
             y      = linear("proj", y);
             x      = ggml_add(c, x, y);
             y      = norm("norm2", x);
-            y      = situ(c, linear("ffn_up", y), linear("ffn_gate", y));
+            y      = ggml_ext_situ(ctx, linear("ffn_up", y), linear("ffn_gate", y));
             y      = linear("ffn_down", y);
             return ggml_add(c, x, y);
         }
@@ -137,7 +131,7 @@ namespace MOSSImage {
             y            = norm("norm2", x);
             y            = ggml_add(c, ggml_add(c, y, ggml_mul(c, y, part(4))), part(3));
             auto up_gate = linear("proj.ffn_up_gate", y);
-            y            = situ(c, ggml_ext_slice(c, up_gate, 0, 0, config.ffn_dim), ggml_ext_slice(c, up_gate, 0, config.ffn_dim, 2 * config.ffn_dim));
+            y            = ggml_ext_situ(ctx, up_gate);
             y            = linear("proj.ffn_down", y);
             y            = ggml_mul(c, y, part(5));
             return ggml_add(c, x, y);