perf: view qkv slices for moss instead of slice+reshape (#2232)

This commit is contained in:
Piotr Wilkin (ilintar) authored and GitHub committed 2026-11-03 23:46:35 +08:00
1 parent f8ea2e162c
commit 6225bf8c88
1 file changed
+2 -1
+2 -1
View File
@@ -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));