From 56070ea0342e0bb93da55c858e952da4543110f9 Mon Sep 17 00:00:00 2001 From: leejet Date: Tue, 10 Nov 2026 23:38:47 +0800 Subject: [PATCH] refactor: replace QKV split helpers with ggml_ext_chunk (#2322) --- src/core/ggml_extend.cpp | 30 ------------------------------ src/core/ggml_extend.h | 8 -------- src/model/common/ggml_block.hpp | 2 +- src/model/diffusion/hunyuan.hpp | 2 +- src/model/diffusion/ideogram4.hpp | 2 +- src/model/diffusion/iris.hpp | 2 +- src/model/diffusion/lens.hpp | 4 ++-- src/model/diffusion/mmdit.hpp | 2 +- src/model/diffusion/pid.hpp | 6 +++--- src/model/te/llm.hpp | 2 +- src/model/vae/wan_vae.hpp | 2 +- 11 files changed, 12 insertions(+), 50 deletions(-) diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 97561581..52d9cfc8 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -541,36 +541,6 @@ ggml_tensor* ggml_ext_conv_3d(GGMLRunnerContext* runner_ctx, return x; } -std::vector ggml_ext_split_qkv(GGMLRunnerContext* runner_ctx, ggml_tensor* qkv) { - auto ctx = runner_ctx->ggml_ctx; - qkv = ggml_reshape_4d(ctx, qkv, qkv->ne[0] / 3, 3, qkv->ne[1], qkv->ne[2]); // [N, L, 3, C] - qkv = ggml_cont(ctx, ggml_permute(ctx, qkv, 0, 3, 1, 2)); // [3, N, L, C] - - int64_t offset = qkv->nb[2] * qkv->ne[2]; - auto q = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 0); // [N, L, C] - auto k = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 1); // [N, L, C] - auto v = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 2); // [N, L, C] - return {q, k, v}; -} - -std::vector ggml_ext_split_image_qkv(GGMLRunnerContext* runner_ctx, ggml_tensor* qkv) { - auto ctx = runner_ctx->ggml_ctx; - int64_t W = qkv->ne[0]; - int64_t H = qkv->ne[1]; - int64_t C = qkv->ne[2] / 3; - int64_t N = qkv->ne[3]; - int64_t nb1 = qkv->nb[1]; - int64_t nb2 = qkv->nb[2]; - qkv = ggml_reshape_4d(ctx, qkv, W * H, C, 3, N); // [N, 3, C, H*W] - qkv = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, qkv, 0, 1, 3, 2)); // [3, N, C, H*W] - - int64_t offset = qkv->nb[2] * qkv->ne[2]; - auto q = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[2], offset * 0); // [N, C, H, W] - auto k = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[2], offset * 1); // [N, C, H, W] - auto v = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[2], offset * 2); // [N, C, H, W] - return {q, k, v}; -} - ggml_tensor* ggml_ext_full(GGMLRunnerContext* runner_ctx, float value, int64_t ne0, diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index f8c14c0e..1f0a42c2 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -150,14 +150,6 @@ ggml_tensor* ggml_ext_conv_3d(GGMLRunnerContext* runner_ctx, float scale = 1.f, bool allow_conv2d = false); -// qkv: [N, L, 3*C] -// return: ([N, L, C], [N, L, C], [N, L, C]) -std::vector ggml_ext_split_qkv(GGMLRunnerContext* runner_ctx, ggml_tensor* qkv); - -// qkv: [N, 3*C, H, W] -// return: ([N, C, H, W], [N, C, H, W], [N, C, H, W]) -std::vector ggml_ext_split_image_qkv(GGMLRunnerContext* runner_ctx, ggml_tensor* qkv); - // Constant and cast helpers require the built-in tensors initialized by GGMLRunner. ggml_tensor* ggml_ext_full(GGMLRunnerContext* runner_ctx, float value, diff --git a/src/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index eb437e5d..805b8682 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -910,7 +910,7 @@ public: if (proj_in) { auto in_proj = std::dynamic_pointer_cast(blocks[in_proj_name]); auto qkv = in_proj->forward(ctx, x); - auto qkv_vec = ggml_ext_split_qkv(ctx, qkv); + auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0); q = qkv_vec[0]; k = qkv_vec[1]; v = qkv_vec[2]; diff --git a/src/model/diffusion/hunyuan.hpp b/src/model/diffusion/hunyuan.hpp index d2bc8b58..15c383de 100644 --- a/src/model/diffusion/hunyuan.hpp +++ b/src/model/diffusion/hunyuan.hpp @@ -50,7 +50,7 @@ namespace Hunyuan { // self attn auto qkv = self_attn_qkv->forward(ctx, norm1->forward(ctx, txt)); - auto qkv_vec = ggml_ext_split_qkv(ctx, qkv); + auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0); auto q = qkv_vec[0]; auto k = qkv_vec[1]; auto v = qkv_vec[2]; diff --git a/src/model/diffusion/ideogram4.hpp b/src/model/diffusion/ideogram4.hpp index fbc189aa..8d7cf2ae 100644 --- a/src/model/diffusion/ideogram4.hpp +++ b/src/model/diffusion/ideogram4.hpp @@ -222,7 +222,7 @@ namespace Ideogram4 { auto out_proj = std::dynamic_pointer_cast(blocks["o"]); auto qkv = qkv_proj->forward(ctx, x); - auto qkv_vec = ggml_ext_split_qkv(ctx, qkv); + 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); diff --git a/src/model/diffusion/iris.hpp b/src/model/diffusion/iris.hpp index 92f4b2e2..662704e2 100644 --- a/src/model/diffusion/iris.hpp +++ b/src/model/diffusion/iris.hpp @@ -92,7 +92,7 @@ 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_split_qkv(ctx, std::dynamic_pointer_cast(blocks["qkv"])->forward(ctx, x)); + 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]); diff --git a/src/model/diffusion/lens.hpp b/src/model/diffusion/lens.hpp index 5c86e65d..92e81d31 100644 --- a/src/model/diffusion/lens.hpp +++ b/src/model/diffusion/lens.hpp @@ -151,8 +151,8 @@ namespace Lens { int64_t N = img->ne[2]; int64_t inner = dim_head * num_heads; - auto img_qkv_vec = ggml_ext_split_qkv(ctx, img_qkv->forward(ctx, img)); - auto txt_qkv_vec = ggml_ext_split_qkv(ctx, txt_qkv->forward(ctx, txt)); + 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); diff --git a/src/model/diffusion/mmdit.hpp b/src/model/diffusion/mmdit.hpp index e409a01f..ea653273 100644 --- a/src/model/diffusion/mmdit.hpp +++ b/src/model/diffusion/mmdit.hpp @@ -334,7 +334,7 @@ public: auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); auto qkv = qkv_proj->forward(ctx, x); - auto qkv_vec = ggml_ext_split_qkv(ctx, qkv); + 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] diff --git a/src/model/diffusion/pid.hpp b/src/model/diffusion/pid.hpp index 225a30ae..372e1572 100644 --- a/src/model/diffusion/pid.hpp +++ b/src/model/diffusion/pid.hpp @@ -293,7 +293,7 @@ namespace Pid { auto proj = std::dynamic_pointer_cast(blocks["proj"]); auto qkv = qkv_proj->forward(ctx, x); - auto qkv_vec = ggml_ext_split_qkv(ctx, qkv); + 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; @@ -343,14 +343,14 @@ namespace Pid { int64_t N = x->ne[2]; int64_t head_dim = dim / num_heads; - auto qkv_x = ggml_ext_split_qkv(ctx, qkv_x_proj->forward(ctx, x)); + 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); qx = q_norm_x->forward(ctx, qx); kx = k_norm_x->forward(ctx, kx); - auto qkv_y = ggml_ext_split_qkv(ctx, qkv_y_proj->forward(ctx, y)); + 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); diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp index 7ce1f175..e6d56260 100644 --- a/src/model/te/llm.hpp +++ b/src/model/te/llm.hpp @@ -1252,7 +1252,7 @@ namespace LLM { } else { auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); auto qkv = qkv_proj->forward(ctx, x); - qkv_vec = ggml_ext_split_qkv(ctx, qkv); + qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0); } 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] diff --git a/src/model/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp index 52b09bd7..d44fe92b 100644 --- a/src/model/vae/wan_vae.hpp +++ b/src/model/vae/wan_vae.hpp @@ -623,7 +623,7 @@ namespace WAN { const int64_t w = x->ne[0]; auto qkv = to_qkv->forward(ctx, x); - auto qkv_vec = ggml_ext_split_image_qkv(ctx, qkv); + auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 2); auto q = qkv_vec[0]; q = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, q, 2, 0, 1, 3)); // [t, h, w, c]