Commit 56070ea for stable-diffusion.cpp
commit 56070ea0342e0bb93da55c858e952da4543110f9
Author: leejet <leejet714@gmail.com>
Date: Tue Nov 10 23:38:47 2026 +0800
refactor: replace QKV split helpers with ggml_ext_chunk (#2322)
diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp
index 9756158..52d9cfc 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_tensor*> 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_tensor*> 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 f8c14c0..1f0a42c 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_tensor*> 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_tensor*> 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 eb437e5..805b868 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<Linear>(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 d2bc8b5..15c383d 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 fbc189a..8d7cf2a 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<Linear>(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 92f4b2e..662704e 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<Linear>(blocks["qkv"])->forward(ctx, x));
+ 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]);
diff --git a/src/model/diffusion/lens.hpp b/src/model/diffusion/lens.hpp
index 5c86e65..92e81d3 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 e409a01..ea65327 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<Linear>(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 225a30a..372e157 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<Linear>(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 7ce1f17..e6d5626 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<Linear>(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 52b09bd..d44fe92 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]