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

refactor: replace QKV split helpers with ggml_ext_chunk (#2322)

Browse Source
This commit is contained in:
leejet authored and GitHub committed 2026-11-10 23:38:47 +08:00
1 parent 0bb132831f
commit 56070ea034
11 files changed
+12 -50
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/core/ggml_extend.cpp
-30
View File
@@ -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,
src/core/ggml_extend.h
-8
View File
@@ -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,
src/model/common/ggml_block.hpp
+1 -1
View File
@@ -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];
src/model/diffusion/hunyuan.hpp
+1 -1
View File
@@ -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];
src/model/diffusion/ideogram4.hpp
+1 -1
View File
@@ -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);
src/model/diffusion/iris.hpp
+1 -1
View File
@@ -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]);
src/model/diffusion/lens.hpp
+2 -2
View File
@@ -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);
src/model/diffusion/mmdit.hpp
+1 -1
View File
@@ -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]
src/model/diffusion/pid.hpp
+3 -3
View File
@@ -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);
src/model/te/llm.hpp
+1 -1
View File
@@ -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]
src/model/vae/wan_vae.hpp
+1 -1
View File
@@ -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]