DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-14 12:56:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf: replace qkv materialized chunks with views (#2335)

Browse Source
This commit is contained in:
Wagner Bruna authored and GitHub committed 2026-11-12 23:10:48 +08:00
1 parent 32d516c019
commit 766ce0ec4e
8 files changed
+123 -91
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/model/diffusion/fibo2.hpp
+24 -20
View File
@@ -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);
};
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];
} else {
q = 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];
} else {
k = project("to_k", kv);
v = project("to_v", kv);
}
}
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]);
};
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);
ggml_tensor *q, *k, *v;
if (fused_qkv) {
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 = reshape(project("to_q", x));
if (fused_kv) {
auto kv_out = project("to_kv", kv);
k = split_views(kv_out, 2, 0);
v = split_views(kv_out, 2, 1);
} else {
k = reshape(project("to_k", kv));
v = reshape(project("to_v", kv));
}
}
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);
src/model/diffusion/ideogram4.hpp
+8 -7
View File
@@ -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_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);
src/model/diffusion/iris.hpp
+13 -8
View File
@@ -92,10 +92,16 @@ 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 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);
}
src/model/diffusion/lens.hpp
+13 -11
View File
@@ -151,20 +151,22 @@ 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);
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 = 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);
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);
src/model/diffusion/mmdit.hpp
+21 -14
View File
@@ -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_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,
src/model/diffusion/pid.hpp
+22 -15
View File
@@ -293,13 +293,15 @@ 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);
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);
@@ -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);
src/model/diffusion/seedvr2.hpp
+10 -3
View File
@@ -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);
src/model/te/llm.hpp
+12 -13
View File
@@ -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"]);
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_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]