diff --git a/src/model/diffusion/moss_image.hpp b/src/model/diffusion/moss_image.hpp index 36d8d8328..1f5dc1668 100644 --- a/src/model/diffusion/moss_image.hpp +++ b/src/model/diffusion/moss_image.hpp @@ -79,7 +79,8 @@ namespace MOSSImage { auto y = norm("norm1", x); auto qkv = linear("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]); + // q/k feed rms_norm and v a permuted copy, which read these strided views in place. + 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));