From 627e126447444df344d4fac7a7ca2c28e8115d7f Mon Sep 17 00:00:00 2001 From: mirek190 Date: Sat, 24 Oct 2026 21:50:51 +0100 Subject: [PATCH 1/4] perf(moss): avoid redundant CUDA projection copies --- docs/moss_image.md | 4 ++++ src/core/ggml_extend.cpp | 13 +++++++++++-- src/model/diffusion/moss_image.hpp | 7 ++++++- 3 files changed, 21 insertions(+), 3 deletions(-) diff --git a/docs/moss_image.md b/docs/moss_image.md index 27ab26f804..3adf26478f 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. + +## CUDA projection layout + +On supported CUDA devices, MOSS consumes Q/K views of the combined projection. For eligible 128-wide attention, the shared attention helper combines V layout materialization and FP16 conversion. Existing paths remain for unsupported layouts, head padding and K/V scaling. This application change works with the pinned ggml; matching backend fusions can provide an additional benefit. 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..9847caf524 100644 --- a/src/model/diffusion/moss_image.hpp +++ b/src/model/diffusion/moss_image.hpp @@ -121,7 +121,12 @@ namespace MOSSImage { 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)); From 042de506b9297d5c6119b9567369ebd35ef09a5d Mon Sep 17 00:00:00 2001 From: mirek190 Date: Sat, 24 Oct 2026 22:44:11 +0100 Subject: [PATCH 2/4] 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 3adf26478f..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. - -## CUDA projection layout - -On supported CUDA devices, MOSS consumes Q/K views of the combined projection. For eligible 128-wide attention, the shared attention helper combines V layout materialization and FP16 conversion. Existing paths remain for unsupported layouts, head padding and K/V scaling. This application change works with the pinned ggml; matching backend fusions can provide an additional benefit. From d8196821b725988a1a5aab59d1d2badbafec7e86 Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 25 Oct 2026 16:51:44 +0800 Subject: [PATCH 3/4] refactor: simplify MOSS QKV views and value preparation --- src/core/ggml_extend.cpp | 21 +++++++-------------- src/model/diffusion/moss_image.hpp | 10 +++------- 2 files changed, 10 insertions(+), 21 deletions(-) diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 574c217ff3..07ecd82342 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -678,30 +678,23 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, // CUDA FA MMA starts at 64 channels; keep the original head's attention scale. q_in = ggml_pad(ctx, q_in, static_cast(64 - d_head), 0, 0, 0); k_in = ggml_pad(ctx, k_in, static_cast(64 - d_head), 0, 0, 0); + v_in = ggml_pad(ctx, v_in, static_cast(64 - d_head), 0, 0, 0); } if (kv_scale != 1.0f) { k_in = ggml_ext_scale(ctx, k_in, kv_scale); } k_in = ggml_cast(ctx, k_in, GGML_TYPE_F16); - 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); - } if (kv_scale != 1.0f) { v_in = ggml_ext_scale(ctx, v_in, kv_scale); } - if (!direct_v_cast) { - v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16); + v_in = ggml_permute(ctx, v_in, 0, 2, 1, 3); + auto v_f16 = ggml_cast(ctx, v_in, GGML_TYPE_F16); + if (true) { + v_f16 = ggml_cast(ctx, ggml_ext_cont(ctx, v_in), GGML_TYPE_F16); } + v_in = v_f16; + v_in = ggml_reshape_3d(ctx, v_in, v_in->ne[0], L_k, n_kv_head * N); 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 9847caf524..fd2a9f3df0 100644 --- a/src/model/diffusion/moss_image.hpp +++ b/src/model/diffusion/moss_image.hpp @@ -121,13 +121,9 @@ namespace MOSSImage { 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); - 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 split = [&](int i) { + 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]); }; auto q = norm("q_norm", split(0)); auto k = norm("k_norm", split(1)); From fdaf518c3e623e2c9f8abd8783579a16b2546bc2 Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 25 Oct 2026 16:53:47 +0800 Subject: [PATCH 4/4] fix: restore capability check for direct value cast --- src/core/ggml_extend.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 07ecd82342..bd5f6f63d0 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -690,7 +690,7 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, } v_in = ggml_permute(ctx, v_in, 0, 2, 1, 3); auto v_f16 = ggml_cast(ctx, v_in, GGML_TYPE_F16); - if (true) { + if (!ggml_backend_supports_op(backend, v_f16)) { v_f16 = ggml_cast(ctx, ggml_ext_cont(ctx, v_in), GGML_TYPE_F16); } v_in = v_f16;