From f7877e77dc38742f357d81254f46c4545d883840 Mon Sep 17 00:00:00 2001 From: mirek190 Date: Sat, 24 Oct 2026 21:50:50 +0100 Subject: [PATCH 1/2] perf(moss): reuse shared modulation slices --- docs/moss_image.md | 4 ++++ src/model/diffusion/moss_image.hpp | 16 ++++++++++++---- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/docs/moss_image.md b/docs/moss_image.md index 27ab26f804..6e278d9762 100644 --- a/docs/moss_image.md +++ b/docs/moss_image.md @@ -35,3 +35,7 @@ Pass the external `tokenizer.json` with `--tokenizer`; Qwen3.5 uses a different Use dimensions divisible by 16. The default discrete scheduler uses the upstream time shift of 5. Use `--flow-shift` to override it. The negative CFG branch skips DiT layer 22 (zero-based), matching the upstream pipeline. + +## Shared modulation + +The model constructs the six shared modulation slices once per forward graph and reuses them across diffusion blocks. Their values and the arithmetic consuming them are unchanged. This application change works with the pinned ggml. diff --git a/src/model/diffusion/moss_image.hpp b/src/model/diffusion/moss_image.hpp index f1cbbc7f37..bea3aa5ede 100644 --- a/src/model/diffusion/moss_image.hpp +++ b/src/model/diffusion/moss_image.hpp @@ -2,6 +2,7 @@ #define __SD_MODEL_DIFFUSION_MOSS_IMAGE_HPP__ #include +#include #include #include "model/common/ggml_block.hpp" @@ -113,11 +114,11 @@ namespace MOSSImage { blocks["proj.ffn_down"] = std::make_shared(config.ffn_dim, config.dim, false); } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, ggml_tensor* mod) { + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::array& mod) { auto c = ctx->ggml_ctx; auto linear = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast(blocks[name])->forward(ctx, h); }; auto norm = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast(blocks[name])->forward(ctx, h); }; - auto part = [&](int i) { return ggml_ext_slice(c, mod, 0, i * config.dim, (i + 1) * config.dim); }; + auto part = [&](int i) { return mod[i]; }; auto y = norm("norm1", x); y = ggml_add(c, ggml_add(c, y, ggml_mul(c, y, part(1))), part(0)); auto qkv = linear("proj.qkv", y); @@ -175,10 +176,17 @@ namespace MOSSImage { time = linear("time_embedding.2", ggml_silu(c, linear("time_embedding.0", time))); auto mod = linear("shared_block_modulation", time); auto head = linear("video_head_modulation", time); + std::array parts; + for (int i = 0; i < 6; ++i) { + parts[i] = ggml_ext_slice(c, mod, 0, i * config.dim, (i + 1) * config.dim); + if (ctx->graph != nullptr) { + ggml_build_forward_expand(ctx->graph, parts[i]); + } + } x = ggml_concat(c, x, text, 1); for (int64_t i = 0; i < config.depth; ++i) { if (std::find(skip.begin(), skip.end(), i) == skip.end()) { - x = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)])->forward(ctx, x, pe, mod); + x = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)])->forward(ctx, x, pe, parts); sd::ggml_graph_cut::mark_graph_cut(x, "moss_image.layers." + std::to_string(i), "x"); } } @@ -250,7 +258,7 @@ namespace MOSSImage { auto time = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.freq_dim, 1, 1); set_backend_tensor_data(pe, pe_data.data()); set_backend_tensor_data(time, time_data.data()); - auto ctx = get_context(); + auto ctx = get_context(gf); ggml_build_forward_expand(gf, model.forward(&ctx, x, time, text, pe, skip)); return gf; }; From dd88193a37301f08196904ebedc370f7198044e3 Mon Sep 17 00:00:00 2001 From: mirek190 Date: Sat, 24 Oct 2026 22:44:10 +0100 Subject: [PATCH 2/2] docs(moss): remove unnecessary optimization documentation --- docs/moss_image.md | 4 ---- 1 file changed, 4 deletions(-) diff --git a/docs/moss_image.md b/docs/moss_image.md index 6e278d9762..27ab26f804 100644 --- a/docs/moss_image.md +++ b/docs/moss_image.md @@ -35,7 +35,3 @@ Pass the external `tokenizer.json` with `--tokenizer`; Qwen3.5 uses a different Use dimensions divisible by 16. The default discrete scheduler uses the upstream time shift of 5. Use `--flow-shift` to override it. The negative CFG branch skips DiT layer 22 (zero-based), matching the upstream pipeline. - -## Shared modulation - -The model constructs the six shared modulation slices once per forward graph and reuses them across diffusion blocks. Their values and the arithmetic consuming them are unchanged. This application change works with the pinned ggml.