Commit 766ce0e for stable-diffusion.cpp
commit 766ce0ec4e3d00c3da3fc6304d4c64fd0cfdab40
Author: Wagner Bruna <wbruna@users.noreply.github.com>
Date: Thu Nov 12 12:10:48 2026 -0300
perf: replace qkv materialized chunks with views (#2335)
diff --git a/src/model/diffusion/fibo2.hpp b/src/model/diffusion/fibo2.hpp
index 1ac7684..6f8ef09 100644
--- a/src/model/diffusion/fibo2.hpp
+++ b/src/model/diffusion/fibo2.hpp
@@ -106,29 +106,33 @@ namespace Fibo2 {
auto project = [&](const char* name, ggml_tensor* h) {
return std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, h);
};
+ auto split_views = [&](ggml_tensor* qkv, int n, int i) {
+ int64_t slice = qkv->ne[0] / n;
+ return ggml_view_4d(c, qkv, head_dim, slice / head_dim, qkv->ne[1], qkv->ne[2] * qkv->ne[3],
+ head_dim * qkv->nb[0], qkv->nb[1], qkv->nb[2], i * slice * qkv->nb[0]);
+ };
+ auto reshape = [&](ggml_tensor* h) {
+ return ggml_reshape_4d(c, h, head_dim, heads, h->ne[1], h->ne[2]);
+ };
ggml_tensor *q, *k, *v;
if (fused_qkv) {
- auto parts = ggml_ext_chunk(ctx, project("qkv", x), 3, 0);
- q = parts[0];
- k = parts[1];
- v = parts[2];
+ auto qkv = project("qkv", x);
+ q = split_views(qkv, 3, 0);
+ k = split_views(qkv, 3, 1);
+ v = split_views(qkv, 3, 2);
} else {
- q = project("to_q", x);
+ q = reshape(project("to_q", x));
if (fused_kv) {
- auto parts = ggml_ext_chunk(ctx, project("to_kv", kv), 2, 0);
- k = parts[0];
- v = parts[1];
+ auto kv_out = project("to_kv", kv);
+ k = split_views(kv_out, 2, 0);
+ v = split_views(kv_out, 2, 1);
} else {
- k = project("to_k", kv);
- v = project("to_v", kv);
+ k = reshape(project("to_k", kv));
+ v = reshape(project("to_v", kv));
}
}
- auto reshape = [&](ggml_tensor* h) {
- return ggml_reshape_4d(c, h, head_dim, heads, h->ne[1], h->ne[2]);
- };
- q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"])->forward(ctx, reshape(q));
- k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"])->forward(ctx, reshape(k));
- v = reshape(v);
+ q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"])->forward(ctx, q);
+ k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"])->forward(ctx, k);
auto prepare = [&](ggml_tensor* h) {
if (pe) {
return Rope::apply_rope(ctx, h, pe);
diff --git a/src/model/diffusion/ideogram4.hpp b/src/model/diffusion/ideogram4.hpp
index 8d7cf2a..bd1abeb 100644
--- a/src/model/diffusion/ideogram4.hpp
+++ b/src/model/diffusion/ideogram4.hpp
@@ -213,19 +213,20 @@ namespace Ideogram4 {
ggml_tensor* x,
ggml_tensor* pe,
ggml_tensor* mask = nullptr) {
- int64_t n_token = x->ne[1];
- int64_t N = x->ne[2];
-
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
auto norm_q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"]);
auto norm_k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["o"]);
- auto qkv = qkv_proj->forward(ctx, x);
- auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0);
- auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, n_token, N);
- auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, n_token, N);
- auto v = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[2], head_dim, num_heads, n_token, N);
+ auto qkv = qkv_proj->forward(ctx, x);
+ auto qkv_dim = qkv->ne[0] / 3;
+ auto split_qkv = [&](int i) {
+ return ggml_view_4d(ctx->ggml_ctx, qkv, head_dim, qkv_dim / head_dim, qkv->ne[1], qkv->ne[2] * qkv->ne[3],
+ head_dim * qkv->nb[0], qkv->nb[1], qkv->nb[2], i * qkv_dim * qkv->nb[0]);
+ };
+ auto q = split_qkv(0);
+ auto k = split_qkv(1);
+ auto v = split_qkv(2);
q = norm_q->forward(ctx, q);
k = norm_k->forward(ctx, k);
diff --git a/src/model/diffusion/iris.hpp b/src/model/diffusion/iris.hpp
index 662704e..eb4c66b 100644
--- a/src/model/diffusion/iris.hpp
+++ b/src/model/diffusion/iris.hpp
@@ -91,11 +91,17 @@ namespace Iris {
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe = nullptr, ggml_tensor* mask = nullptr) {
- auto c = ctx->ggml_ctx;
- auto qkv = ggml_ext_chunk(ctx, std::dynamic_pointer_cast<Linear>(blocks["qkv"])->forward(ctx, x), 3, 0);
- auto q = ggml_reshape_4d(c, qkv[0], dim / heads, heads, x->ne[1], x->ne[2]);
- auto k = ggml_reshape_4d(c, qkv[1], dim / heads, heads, x->ne[1], x->ne[2]);
- auto v = ggml_reshape_4d(c, qkv[2], dim / heads, heads, x->ne[1], x->ne[2]);
+ auto c = ctx->ggml_ctx;
+ auto head_dim = dim / heads;
+ auto qkv = std::dynamic_pointer_cast<Linear>(blocks["qkv"])->forward(ctx, x);
+ auto qkv_dim = qkv->ne[0] / 3;
+ auto split_qkv = [&](int i) {
+ return ggml_view_4d(c, qkv, head_dim, qkv_dim / head_dim, qkv->ne[1], qkv->ne[2] * qkv->ne[3],
+ head_dim * qkv->nb[0], qkv->nb[1], qkv->nb[2], i * qkv_dim * qkv->nb[0]);
+ };
+ auto q = split_qkv(0);
+ auto k = split_qkv(1);
+ auto v = split_qkv(2);
if (qk_norm) {
q = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"])->forward(ctx, q);
k = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"])->forward(ctx, k);
@@ -103,10 +109,9 @@ namespace Iris {
if (pe) {
x = Rope::attention(ctx, q, k, v, pe, mask);
} else {
- q = ggml_reshape_3d(c, q, dim, x->ne[1], x->ne[2]);
- k = ggml_reshape_3d(c, k, dim, x->ne[1], x->ne[2]);
- v = ggml_reshape_3d(c, v, dim, x->ne[1], x->ne[2]);
- x = ggml_ext_attention_ext(ctx, q, k, v, heads, mask, false, ctx->flash_attn_enabled);
+ q = ggml_reshape_3d(c, ggml_cont(c, ggml_permute(c, q, 0, 2, 1, 3)), head_dim, x->ne[1], heads * x->ne[2]);
+ k = ggml_reshape_3d(c, ggml_cont(c, ggml_permute(c, k, 0, 2, 1, 3)), head_dim, x->ne[1], heads * x->ne[2]);
+ x = ggml_ext_attention_ext(ctx, q, k, v, heads, mask, true, ctx->flash_attn_enabled);
}
return std::dynamic_pointer_cast<Linear>(blocks["proj"])->forward(ctx, x);
}
diff --git a/src/model/diffusion/lens.hpp b/src/model/diffusion/lens.hpp
index 92e81d3..85c987f 100644
--- a/src/model/diffusion/lens.hpp
+++ b/src/model/diffusion/lens.hpp
@@ -151,22 +151,24 @@ namespace Lens {
int64_t N = img->ne[2];
int64_t inner = dim_head * num_heads;
- auto img_qkv_vec = ggml_ext_chunk(ctx, img_qkv->forward(ctx, img), 3, 0);
- auto txt_qkv_vec = ggml_ext_chunk(ctx, txt_qkv->forward(ctx, txt), 3, 0);
-
- auto img_q = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[0], dim_head, num_heads, n_img, N);
- auto img_k = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[1], dim_head, num_heads, n_img, N);
- auto img_v = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[2], dim_head, num_heads, n_img, N);
-
- img_q = norm_q->forward(ctx, img_q);
- img_k = norm_k->forward(ctx, img_k);
-
- auto txt_q = ggml_reshape_4d(ctx->ggml_ctx, txt_qkv_vec[0], dim_head, num_heads, n_txt, N);
- auto txt_k = ggml_reshape_4d(ctx->ggml_ctx, txt_qkv_vec[1], dim_head, num_heads, n_txt, N);
- auto txt_v = ggml_reshape_4d(ctx->ggml_ctx, txt_qkv_vec[2], dim_head, num_heads, n_txt, N);
-
- txt_q = norm_add_q->forward(ctx, txt_q);
- txt_k = norm_add_k->forward(ctx, txt_k);
+ auto img_qkv_f = img_qkv->forward(ctx, img);
+ auto txt_qkv_f = txt_qkv->forward(ctx, txt);
+ auto split_qkv = [&](ggml_tensor* qkv, int i) {
+ int64_t qkv_dim = qkv->ne[0] / 3;
+ return ggml_view_4d(ctx->ggml_ctx, qkv, dim_head, qkv_dim / dim_head, qkv->ne[1], qkv->ne[2] * qkv->ne[3],
+ dim_head * qkv->nb[0], qkv->nb[1], qkv->nb[2], i * qkv_dim * qkv->nb[0]);
+ };
+ auto img_q = split_qkv(img_qkv_f, 0);
+ auto img_k = split_qkv(img_qkv_f, 1);
+ auto img_v = split_qkv(img_qkv_f, 2);
+ img_q = norm_q->forward(ctx, img_q);
+ img_k = norm_k->forward(ctx, img_k);
+
+ auto txt_q = split_qkv(txt_qkv_f, 0);
+ auto txt_k = split_qkv(txt_qkv_f, 1);
+ auto txt_v = split_qkv(txt_qkv_f, 2);
+ txt_q = norm_add_q->forward(ctx, txt_q);
+ txt_k = norm_add_k->forward(ctx, txt_k);
auto q = ggml_concat(ctx->ggml_ctx, img_q, txt_q, 2);
auto k = ggml_concat(ctx->ggml_ctx, img_k, txt_k, 2);
diff --git a/src/model/diffusion/mmdit.hpp b/src/model/diffusion/mmdit.hpp
index ea65327..1f1fed4 100644
--- a/src/model/diffusion/mmdit.hpp
+++ b/src/model/diffusion/mmdit.hpp
@@ -331,14 +331,19 @@ public:
}
std::vector<ggml_tensor*> pre_attention(GGMLRunnerContext* ctx, ggml_tensor* x) {
+ auto c = ctx->ggml_ctx;
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
- auto qkv = qkv_proj->forward(ctx, x);
- auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0);
- int64_t head_dim = qkv_vec[0]->ne[0] / num_heads;
- auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, qkv_vec[0]->ne[1], qkv_vec[0]->ne[2]); // [N, n_token, n_head, d_head]
- auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, qkv_vec[1]->ne[1], qkv_vec[1]->ne[2]); // [N, n_token, n_head, d_head]
- auto v = qkv_vec[2]; // [N, n_token, n_head*d_head]
+ auto qkv = qkv_proj->forward(ctx, x);
+ auto qkv_dim = qkv->ne[0] / 3;
+ auto head_dim = qkv_dim / num_heads;
+ auto split_qkv = [&](int i) {
+ return ggml_view_4d(c, qkv, head_dim, qkv_dim / head_dim, qkv->ne[1], qkv->ne[2] * qkv->ne[3],
+ head_dim * qkv->nb[0], qkv->nb[1], qkv->nb[2], i * qkv_dim * qkv->nb[0]);
+ };
+ auto q = split_qkv(0);
+ auto k = split_qkv(1);
+ auto v = split_qkv(2);
if (qk_norm == "rms" || qk_norm == "ln") {
auto ln_q = std::dynamic_pointer_cast<UnaryBlock>(blocks["ln_q"]);
@@ -347,8 +352,9 @@ public:
k = ln_k->forward(ctx, k);
}
- q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0] * q->ne[1], q->ne[2], q->ne[3]); // [N, n_token, n_head*d_head]
- k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0] * k->ne[1], k->ne[2], k->ne[3]); // [N, n_token, n_head*d_head]
+ // skip_reshape=true contract: q/k become [d_head, L, n_head*N]; v stays a 4-D view.
+ q = ggml_reshape_3d(c, ggml_cont(c, ggml_permute(c, q, 0, 2, 1, 3)), q->ne[0], q->ne[2], num_heads * q->ne[3]);
+ k = ggml_reshape_3d(c, ggml_cont(c, ggml_permute(c, k, 0, 2, 1, 3)), k->ne[0], k->ne[2], num_heads * k->ne[3]);
return {q, k, v};
}
@@ -366,7 +372,7 @@ public:
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x) {
auto qkv = pre_attention(ctx, x);
- x = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
+ x = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, true, ctx->flash_attn_enabled); // [N, n_token, dim]
x = post_attention(ctx, x); // [N, n_token, dim]
return x;
}
@@ -586,8 +592,8 @@ public:
auto qkv2 = std::get<1>(qkv_intermediates);
auto intermediates = std::get<2>(qkv_intermediates);
- auto attn_out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
- auto attn2_out = ggml_ext_attention_ext(ctx, qkv2[0], qkv2[1], qkv2[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
+ auto attn_out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, true, ctx->flash_attn_enabled); // [N, n_token, dim]
+ auto attn2_out = ggml_ext_attention_ext(ctx, qkv2[0], qkv2[1], qkv2[2], num_heads, nullptr, true, ctx->flash_attn_enabled); // [N, n_token, dim]
x = post_attention_x(ctx,
attn_out,
attn2_out,
@@ -603,7 +609,7 @@ public:
auto qkv = qkv_intermediates.first;
auto intermediates = qkv_intermediates.second;
- auto attn_out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
+ auto attn_out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, true, ctx->flash_attn_enabled); // [N, n_token, dim]
x = post_attention(ctx,
attn_out,
intermediates[0],
@@ -644,10 +650,11 @@ block_mixing(GGMLRunnerContext* ctx,
}
std::vector<ggml_tensor*> qkv;
for (int i = 0; i < 3; i++) {
- qkv.push_back(ggml_concat(ctx->ggml_ctx, context_qkv[i], x_qkv[i], 1));
+ // q/k are 3-D [dim, L, N] (token axis = 1); v is a 4-D view [d_head, heads, L, N] (token axis = 2)
+ qkv.push_back(ggml_concat(ctx->ggml_ctx, context_qkv[i], x_qkv[i], i == 2 ? 2 : 1));
}
- auto attn = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_context + n_token, hidden_size]
+ auto attn = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], x_block->num_heads, nullptr, true, ctx->flash_attn_enabled); // [N, n_context + n_token, hidden_size]
auto context_attn = ggml_view_3d(ctx->ggml_ctx,
attn,
@@ -679,7 +686,7 @@ block_mixing(GGMLRunnerContext* ctx,
}
if (x_block->self_attn) {
- auto attn2 = ggml_ext_attention_ext(ctx, x_qkv2[0], x_qkv2[1], x_qkv2[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, hidden_size]
+ auto attn2 = ggml_ext_attention_ext(ctx, x_qkv2[0], x_qkv2[1], x_qkv2[2], x_block->num_heads, nullptr, true, ctx->flash_attn_enabled); // [N, n_token, hidden_size]
x = x_block->post_attention_x(ctx,
x_attn,
diff --git a/src/model/diffusion/pid.hpp b/src/model/diffusion/pid.hpp
index 372e157..1002bf7 100644
--- a/src/model/diffusion/pid.hpp
+++ b/src/model/diffusion/pid.hpp
@@ -293,16 +293,18 @@ namespace Pid {
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
auto qkv = qkv_proj->forward(ctx, x);
- auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0);
- int64_t L = x->ne[1];
- int64_t N = x->ne[2];
int64_t head_dim = dim / num_heads;
- auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, L, N);
- auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, L, N);
- auto v = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[2], head_dim, num_heads, L, N);
- q = q_norm->forward(ctx, q);
- k = k_norm->forward(ctx, k);
- x = Rope::attention(ctx, q, k, v, pos, nullptr, 1.0f / 128.f, true);
+ auto qkv_dim = qkv->ne[0] / 3;
+ auto split_qkv = [&](int i) {
+ return ggml_view_4d(ctx->ggml_ctx, qkv, head_dim, qkv_dim / head_dim, qkv->ne[1], qkv->ne[2] * qkv->ne[3],
+ head_dim * qkv->nb[0], qkv->nb[1], qkv->nb[2], i * qkv_dim * qkv->nb[0]);
+ };
+ auto q = split_qkv(0);
+ auto k = split_qkv(1);
+ auto v = split_qkv(2);
+ q = q_norm->forward(ctx, q);
+ k = k_norm->forward(ctx, k);
+ x = Rope::attention(ctx, q, k, v, pos, nullptr, 1.0f / 128.f, true);
return proj->forward(ctx, x);
}
};
@@ -340,20 +342,25 @@ namespace Pid {
int64_t Nx = x->ne[1];
int64_t Ny = y->ne[1];
- int64_t N = x->ne[2];
int64_t head_dim = dim / num_heads;
- auto qkv_x = ggml_ext_chunk(ctx, qkv_x_proj->forward(ctx, x), 3, 0);
- auto qx = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[0], head_dim, num_heads, Nx, N);
- auto kx = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[1], head_dim, num_heads, Nx, N);
- auto vx = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[2], head_dim, num_heads, Nx, N);
+ auto split_qkv = [&](ggml_tensor* qkv, int i) {
+ int64_t qkv_dim = qkv->ne[0] / 3;
+ return ggml_view_4d(ctx->ggml_ctx, qkv, head_dim, qkv_dim / head_dim, qkv->ne[1], qkv->ne[2] * qkv->ne[3],
+ head_dim * qkv->nb[0], qkv->nb[1], qkv->nb[2], i * qkv_dim * qkv->nb[0]);
+ };
+
+ auto qkv_x = qkv_x_proj->forward(ctx, x);
+ auto qx = split_qkv(qkv_x, 0);
+ auto kx = split_qkv(qkv_x, 1);
+ auto vx = split_qkv(qkv_x, 2);
qx = q_norm_x->forward(ctx, qx);
kx = k_norm_x->forward(ctx, kx);
- auto qkv_y = ggml_ext_chunk(ctx, qkv_y_proj->forward(ctx, y), 3, 0);
- auto qy = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[0], head_dim, num_heads, Ny, N);
- auto ky = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[1], head_dim, num_heads, Ny, N);
- auto vy = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[2], head_dim, num_heads, Ny, N);
+ auto qkv_y = qkv_y_proj->forward(ctx, y);
+ auto qy = split_qkv(qkv_y, 0);
+ auto ky = split_qkv(qkv_y, 1);
+ auto vy = split_qkv(qkv_y, 2);
qy = q_norm_y->forward(ctx, qy);
ky = k_norm_y->forward(ctx, ky);
diff --git a/src/model/diffusion/seedvr2.hpp b/src/model/diffusion/seedvr2.hpp
index 0f0a62e..4cb93bd 100644
--- a/src/model/diffusion/seedvr2.hpp
+++ b/src/model/diffusion/seedvr2.hpp
@@ -199,14 +199,21 @@ namespace SeedVR2 {
vqkv = ggml_get_rows(c, vqkv, inputs.input(ctx, plan.order));
auto vpe = rotation(ctx, inputs.input(ctx, plan.positions));
auto tpe = rotation(ctx, inputs.input(ctx, plan.text_positions));
+ auto split_qkv = [&](ggml_tensor* x, int i) {
+ int64_t qkv_dim = x->ne[0] / 3;
+ return ggml_view_4d(c, x, 128, qkv_dim / 128, x->ne[1], x->ne[2] * x->ne[3],
+ 128 * x->nb[0], x->nb[1], x->nb[2], i * qkv_dim * x->nb[0]);
+ };
auto qkv = [&](ggml_tensor* x, bool text, ggml_tensor* pe) {
- auto parts = ggml_ext_chunk(ctx, x, 3, 0);
+ std::vector<ggml_tensor*> parts{nullptr, nullptr, nullptr};
for (int i = 0; i < 2; ++i) {
- auto h = ggml_reshape_3d(c, parts[i], 128, heads, x->ne[1]);
- h = std::dynamic_pointer_cast<RMSNorm>(blocks[std::string(i == 0 ? "attn.norm_q." : "attn.norm_k.") + branch(text)])->forward(ctx, h);
+ auto h = std::dynamic_pointer_cast<RMSNorm>(blocks[std::string(i == 0 ? "attn.norm_q." : "attn.norm_k.") + branch(text)])->forward(ctx, split_qkv(x, i));
h = ggml_reshape_2d(c, h, dim, x->ne[1]);
parts[i] = VSR::apply_rope(ctx, h, pe, heads);
}
+ // v must stay 3-D [dim, L, N]: the window loop slices dim 1 as tokens.
+ auto v = split_qkv(x, 2);
+ parts[2] = ggml_view_3d(c, v, dim, x->ne[1], x->ne[2], v->nb[2], v->nb[3], 0);
return parts;
};
auto vp = qkv(vqkv, false, vpe), tp = qkv(tqkv, true, tpe);
diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp
index e6d5626..3bba320 100644
--- a/src/model/te/llm.hpp
+++ b/src/model/te/llm.hpp
@@ -1234,31 +1234,30 @@ namespace LLM {
ggml_tensor* pe,
ggml_tensor* mask = nullptr) {
// x: [N, n_token, hidden_size]
- int64_t n_token = x->ne[1];
- int64_t N = x->ne[2];
- auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
+ auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
- std::vector<ggml_tensor*> qkv_vec;
+ ggml_tensor *q, *k, *v;
if (llama_cpp_style) {
auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["q_proj"]);
auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
auto v_proj = std::dynamic_pointer_cast<Linear>(blocks["v_proj"]);
- auto q = q_proj->forward(ctx, x);
- auto k = k_proj->forward(ctx, x);
- auto v = v_proj->forward(ctx, x);
-
- qkv_vec = {q, k, v};
+ q = ggml_reshape_4d(ctx->ggml_ctx, q_proj->forward(ctx, x), head_dim, num_heads, x->ne[1], x->ne[2]);
+ k = ggml_reshape_4d(ctx->ggml_ctx, k_proj->forward(ctx, x), head_dim, num_heads, x->ne[1], x->ne[2]);
+ v = ggml_reshape_4d(ctx->ggml_ctx, v_proj->forward(ctx, x), head_dim, num_heads, x->ne[1], x->ne[2]);
} else {
- auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
- auto qkv = qkv_proj->forward(ctx, x);
- qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0);
+ auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
+ auto qkv = qkv_proj->forward(ctx, x);
+ auto qkv_dim = qkv->ne[0] / 3;
+ auto split_qkv = [&](int i) {
+ return ggml_view_4d(ctx->ggml_ctx, qkv, head_dim, qkv_dim / head_dim, qkv->ne[1], qkv->ne[2] * qkv->ne[3],
+ head_dim * qkv->nb[0], qkv->nb[1], qkv->nb[2], i * qkv_dim * qkv->nb[0]);
+ };
+ q = split_qkv(0);
+ k = split_qkv(1);
+ v = split_qkv(2);
}
- auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, qkv_vec[0]->ne[1], qkv_vec[0]->ne[2]); // [N, n_token, n_head, d_head]
- auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, qkv_vec[1]->ne[1], qkv_vec[1]->ne[2]); // [N, n_token, n_head, d_head]
- auto v = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[2], head_dim, num_heads, qkv_vec[2]->ne[1], qkv_vec[2]->ne[2]); // [N, n_token, n_head, d_head]
-
x = Rope::attention(ctx, q, k, v, pe, mask, 1.f, false); // [N, n_token, hidden_size]
x = proj->forward(ctx, x); // [N, n_token, hidden_size]