diff --git a/ggml b/ggml index d49c49e1d..ccb88f810 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit d49c49e1dd361ba3c0a77c90b5b1404b83969aa8 +Subproject commit ccb88f810be05bccc9621e706354611ccc2462c7 diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index ca1b231c8..34d35056a 100644 --- a/src/core/ggml_extend.cpp +++ b/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 diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index 9ddfb0ad6..791155780 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -52,6 +52,9 @@ std::vector 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); diff --git a/src/model/diffusion/moss_image.hpp b/src/model/diffusion/moss_image.hpp index 188c74196..36d8d8328 100644 --- a/src/model/diffusion/moss_image.hpp +++ b/src/model/diffusion/moss_image.hpp @@ -4,6 +4,7 @@ #include #include #include + #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);