From e6370bf093c50260190159237ac80954634cdf3b Mon Sep 17 00:00:00 2001 From: Mirek Date: Sun, 25 Oct 2026 08:54:51 +0000 Subject: [PATCH] perf: avoid redundant MOSS QKV and value layout copies (#2209) Co-authored-by: leejet --- src/core/ggml_extend.cpp | 14 ++++++++------ src/model/diffusion/moss_image.hpp | 3 ++- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 1099f4b0..bd5f6f63 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -678,21 +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_ext_cont(ctx, ggml_permute(ctx, v_in, 0, 2, 1, 3)); - 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); } - 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 (!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; + 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 7aaf9f37..a854c327 100644 --- a/src/model/diffusion/moss_image.hpp +++ b/src/model/diffusion/moss_image.hpp @@ -124,7 +124,8 @@ namespace MOSSImage { 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) { - 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]); + 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));