diff --git a/assets/moss/cuda-q8-original-vs-current.png b/assets/moss/cuda-q8-original-vs-current.png new file mode 100644 index 0000000000..f119d817f8 Binary files /dev/null and b/assets/moss/cuda-q8-original-vs-current.png differ diff --git a/docs/moss_image.md b/docs/moss_image.md index 27ab26f804..77b87727be 100644 --- a/docs/moss_image.md +++ b/docs/moss_image.md @@ -35,3 +35,36 @@ 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. + +## CUDA performance + +With the matching ggml backend changes, MOSS can reuse positional embeddings +and modulation slices, consume Q/K views without materializing the combined +projection, and use fused SiTU activation. Eligible SM86 CUDA paths combine +activation or modulated RMS normalization with Q8 matrix-input packing, and +combine head normalization, rotary positioning and the final FP16 K conversion. +The attention path reuses K/V tiles across more queries and overlaps complete +tile loads with computation. Unsupported devices, shapes and layouts retain +their existing paths. SageAttention is a separate option. + +Fused SiTU is compiled only when ggml advertises `GGML_HAS_SITU`; older +headers keep the original composite activation. Backend capability checks +still select the fallback when the operation is unavailable on a device. + +The following saved measurements use an RTX 3090, MOSS-Image-SFT Q8_0, +Qwen3.5-9B Q8_0, the FLUX.2 VAE, 1024x1024, 10 Euler steps, CFG 3, flow shift 5, +seed 42 and automatic VRAM placement. Ordinary Flash Attention is enabled; +SageAttention, CUDA Graphs and the server conditioning cache are disabled. + +| Implementation | Generation | Denoising per step | Peak VRAM | +|---|---:|---:|---:| +| Original | 17.88 s | 1.731 s | 17.04 GiB | +| Optimized | 12.18 s | 1.161 s | 17.04 GiB | + +This is 31.9% less measured generation time. The decoded RGB arrays match +exactly for this configuration. Each build had one warmup and one timed image; +the original and optimized results were recorded in separate sessions. Startup +and HTTP/PNG transfer are excluded, and timing uncertainty was not quantified. +Pixel equality and this speedup are not established for untested configurations. + +![Original and optimized MOSS output with generation times](../assets/moss/cuda-q8-original-vs-current.png) diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 1099f4b037..574c217ff3 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -684,7 +684,14 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, } k_in = ggml_cast(ctx, k_in, GGML_TYPE_F16); - v_in = ggml_ext_cont(ctx, ggml_permute(ctx, v_in, 0, 2, 1, 3)); + v_in = ggml_permute(ctx, v_in, 0, 2, 1, 3); + const bool direct_v_cast = !pad_head && kv_scale == 1.0f && d_head == 128 && + v_in->type == GGML_TYPE_F32 && sd_backend_supports_cuda_mma(backend); + if (direct_v_cast) { + v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16); + } else { + v_in = ggml_ext_cont(ctx, v_in); + } v_in = ggml_reshape_3d(ctx, v_in, d_head, L_k, n_kv_head * N); if (pad_head) { v_in = ggml_pad(ctx, v_in, static_cast(64 - d_head), 0, 0, 0); @@ -692,7 +699,9 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, if (kv_scale != 1.0f) { v_in = ggml_ext_scale(ctx, v_in, kv_scale); } - v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16); + if (!direct_v_cast) { + v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16); + } if (mask_in != nullptr) { // ggml_flash_attn_ext expects the mask as a contiguous F16 tensor shaped diff --git a/src/model/diffusion/moss_image.hpp b/src/model/diffusion/moss_image.hpp index f1cbbc7f37..58b775350c 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" @@ -53,7 +54,14 @@ namespace MOSSImage { } }; - inline ggml_tensor* situ(ggml_context* c, ggml_tensor* hidden, ggml_tensor* gate) { + inline ggml_tensor* situ(GGMLRunnerContext* ctx, ggml_tensor* hidden, ggml_tensor* gate) { + auto c = ctx->ggml_ctx; +#if !defined(SD_USE_UPSTREAM_GGML) && defined(GGML_HAS_SITU) + auto out = ggml_glu_split(c, gate, hidden, GGML_GLU_OP_SITU); + if (ctx->backend != nullptr && ggml_backend_supports_op(ctx->backend, out)) { + return out; + } +#endif 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); @@ -92,7 +100,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 = situ(ctx, linear("ffn_up", y), linear("ffn_gate", y)); y = linear("ffn_down", y); return ggml_add(c, x, y); } @@ -113,15 +121,20 @@ 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); - auto split = [&](int i) { + const bool use_qkv_views = sd_backend_supports_cuda_mma(ctx->backend); + auto split = [&](int i) { + if (use_qkv_views) { + return ggml_view_4d(c, qkv, config.head_dim, config.dim / config.head_dim, x->ne[1], x->ne[2], + config.head_dim * qkv->nb[0], qkv->nb[1], qkv->nb[2], i * config.dim * qkv->nb[0]); + } return ggml_reshape_4d(c, ggml_ext_slice(c, qkv, 0, i * config.dim, (i + 1) * config.dim), config.head_dim, config.dim / config.head_dim, x->ne[1], x->ne[2]); }; auto q = norm("q_norm", split(0)); @@ -134,7 +147,16 @@ 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 = nullptr; +#if !defined(SD_USE_UPSTREAM_GGML) && defined(GGML_HAS_SITU) + auto activation = ggml_glu(c, up_gate, GGML_GLU_OP_SITU, true); + if (ctx->backend != nullptr && ggml_backend_supports_op(ctx->backend, activation)) { + y = activation; + } +#endif + if (y == nullptr) { + y = situ(ctx, 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 = linear("proj.ffn_down", y); y = ggml_mul(c, y, part(5)); return ggml_add(c, x, y); @@ -175,10 +197,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"); } } @@ -193,9 +222,20 @@ namespace MOSSImage { }; struct MOSSImageRunner : public DiffusionModelRunner { + struct PositionCacheEntry { + int64_t height = 0; + int64_t width = 0; + int64_t text_tokens = 0; + bool circular_x = false; + bool circular_y = false; + std::vector values; + }; + MOSSImageConfig config; Model model; - std::vector pe_data, time_data; + std::array position_cache; + size_t next_position_cache_entry = 0; + std::vector time_data; MOSSImageRunner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr weight_manager = nullptr) : DiffusionModelRunner(backend, prefix, weight_manager), config(MOSSImageConfig::detect_from_weights(weights, prefix)), model(config) { @@ -207,6 +247,39 @@ namespace MOSSImage { model.get_param_tensors(tensors, prefix); } + const std::vector& position_embedding(int64_t h, int64_t w, int64_t nt) { + for (const auto& entry : position_cache) { + if (!entry.values.empty() && entry.height == h && entry.width == w && entry.text_tokens == nt && + entry.circular_x == circular_x_enabled && entry.circular_y == circular_y_enabled) { + return entry.values; + } + } + Rope::Embedding embedding; + embedding.positions.append_image(static_cast(h), static_cast(w)); + embedding.positions.images.back().height_axis = 2; + embedding.positions.images.back().width_axis = 3; + embedding.positions.append_tokens(nt); + for (int64_t token = 0; token < h * w + nt; ++token) { + float pos[4] = {0.f, float(token - h * w), 0.f, 0.f}; + if (token < h * w) { + pos[1] = float(nt); + pos[2] = float(token / w) - float(h - 1) / 2.f; + pos[3] = float(token % w) - float(w - 1) / 2.f; + } + embedding.ids.push_back({pos[0], pos[1], pos[2], pos[3]}); + } + embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, std::vector(4, static_cast(config.head_dim / 4)), embedding.layout, &embedding.frequencies); + auto& entry = position_cache[next_position_cache_entry]; + next_position_cache_entry = (next_position_cache_entry + 1) % position_cache.size(); + entry.height = h; + entry.width = w; + entry.text_tokens = nt; + entry.circular_x = circular_x_enabled; + entry.circular_y = circular_y_enabled; + entry.values = finish_rope_pe(std::move(embedding)); + return entry.values; + } + sd::Tensor compute(int n_threads, const DiffusionParams& inputs) override { if (!inputs.x || !inputs.context || !inputs.timesteps || inputs.context->empty()) { LOG_ERROR("MOSS-Image requires latents, timesteps and text features"); @@ -224,22 +297,7 @@ namespace MOSSImage { auto text = make_input(*inputs.context); GGML_ASSERT(x->ne[3] == 1 && x->ne[2] == config.channels && text->ne[0] == config.text_dim); int64_t h = x->ne[1], w = x->ne[0], nt = text->ne[1]; - Rope::Embedding embedding; - embedding.positions.append_image(static_cast(h), static_cast(w)); - embedding.positions.images.back().height_axis = 2; - embedding.positions.images.back().width_axis = 3; - embedding.positions.append_tokens(nt); - for (int64_t token = 0; token < h * w + nt; ++token) { - float pos[4] = {0.f, float(token - h * w), 0.f, 0.f}; - if (token < h * w) { - pos[1] = float(nt); - pos[2] = float(token / w) - float(h - 1) / 2.f; - pos[3] = float(token % w) - float(w - 1) / 2.f; - } - embedding.ids.push_back({pos[0], pos[1], pos[2], pos[3]}); - } - embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, std::vector(4, static_cast(config.head_dim / 4)), embedding.layout, &embedding.frequencies); - pe_data = finish_rope_pe(std::move(embedding)); + const auto& pe_data = position_embedding(h, w, nt); time_data.resize(config.freq_dim); for (int64_t i = 0; i < config.freq_dim / 2; ++i) { double angle = inputs.timesteps->values()[0] * std::pow(10000.0, -double(i) / double(config.freq_dim / 2)); @@ -250,7 +308,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; };