From 766ce0ec4e3d00c3da3fc6304d4c64fd0cfdab40 Mon Sep 17 00:00:00 2001 From: Wagner Bruna Date: Thu, 12 Nov 2026 12:10:48 -0300 Subject: [PATCH] perf: replace qkv materialized chunks with views (#2335) --- src/model/diffusion/fibo2.hpp | 44 +++++++++++++++++-------------- src/model/diffusion/ideogram4.hpp | 17 ++++++------ src/model/diffusion/iris.hpp | 23 +++++++++------- src/model/diffusion/lens.hpp | 32 +++++++++++----------- src/model/diffusion/mmdit.hpp | 37 +++++++++++++++----------- src/model/diffusion/pid.hpp | 43 +++++++++++++++++------------- src/model/diffusion/seedvr2.hpp | 13 ++++++--- src/model/te/llm.hpp | 31 +++++++++++----------- 8 files changed, 136 insertions(+), 104 deletions(-) diff --git a/src/model/diffusion/fibo2.hpp b/src/model/diffusion/fibo2.hpp index 1ac76843..6f8ef09c 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(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(blocks["norm_q"])->forward(ctx, reshape(q)); - k = std::dynamic_pointer_cast(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(blocks["norm_q"])->forward(ctx, q); + k = std::dynamic_pointer_cast(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 8d7cf2ae..bd1abebe 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(blocks["qkv"]); auto norm_q = std::dynamic_pointer_cast(blocks["norm_q"]); auto norm_k = std::dynamic_pointer_cast(blocks["norm_k"]); auto out_proj = std::dynamic_pointer_cast(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 662704e2..eb4c66bf 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(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(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(blocks["q_norm"])->forward(ctx, q); k = std::dynamic_pointer_cast(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(blocks["proj"])->forward(ctx, x); } diff --git a/src/model/diffusion/lens.hpp b/src/model/diffusion/lens.hpp index 92e81d31..85c987fc 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_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 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 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 ea653273..1f1fed4f 100644 --- a/src/model/diffusion/mmdit.hpp +++ b/src/model/diffusion/mmdit.hpp @@ -331,14 +331,19 @@ public: } std::vector pre_attention(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto c = ctx->ggml_ctx; auto qkv_proj = std::dynamic_pointer_cast(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(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 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 372e1572..1002bf73 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(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 0f0a62e9..4cb93bd1 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 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(blocks[std::string(i == 0 ? "attn.norm_q." : "attn.norm_k.") + branch(text)])->forward(ctx, h); + auto h = std::dynamic_pointer_cast(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 e6d56260..3bba3205 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(blocks["proj"]); + auto proj = std::dynamic_pointer_cast(blocks["proj"]); - std::vector qkv_vec; + ggml_tensor *q, *k, *v; if (llama_cpp_style) { auto q_proj = std::dynamic_pointer_cast(blocks["q_proj"]); auto k_proj = std::dynamic_pointer_cast(blocks["k_proj"]); auto v_proj = std::dynamic_pointer_cast(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(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(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]