From 3e037a81e472fb0abb82a06f421a29cbb2d703b1 Mon Sep 17 00:00:00 2001 From: leejet Date: Sat, 19 Sep 2026 17:50:01 +0800 Subject: [PATCH] perf: accelerate VAE direct 3D convolutions (#1996) --- examples/common/common.cpp | 2 +- ggml | 2 +- src/core/ggml_extend.cpp | 9 +++++++-- src/core/ggml_extend.h | 3 ++- src/core/ggml_runner.cpp | 1 + src/core/ggml_runner.h | 6 ++++++ src/model/common/ggml_block.hpp | 2 +- src/model/vae/wan_vae.hpp | 3 ++- src/pipeline/model_builders.cpp | 4 +++- 9 files changed, 24 insertions(+), 8 deletions(-) diff --git a/examples/common/common.cpp b/examples/common/common.cpp index e8ec4029..933c8622 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -624,7 +624,7 @@ ArgOptions SDContextParams::get_options() { true, &diffusion_conv_direct}, {"", "--vae-conv-direct", - "use ggml_conv2d_direct in the vae model", + "use direct 2D and 3D convolutions in the vae model", true, &vae_conv_direct}, }; diff --git a/ggml b/ggml index 1e22ec0d..c6632cd9 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit 1e22ec0d04b43afa69963e4b4ea6683535d54d79 +Subproject commit c6632cd905401abc58b6f5cdd52d228aa7ca1b88 diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 36f70d79..ab55a946 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -452,8 +452,13 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, int d0, int d1, int d2, - bool force_prec_f32) { - if (force_prec_f32) { + bool force_prec_f32, + bool direct) { + if (direct) { + int64_t OC = w->ne[3] / IC; + int64_t N = x->ne[3] / IC; + x = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC); + } else if (force_prec_f32) { ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type); int64_t OC = w->ne[3] / IC; diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index 6f3fe89e..e843cbee 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -153,7 +153,8 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, int d0 = 1, int d1 = 1, int d2 = 1, - bool force_prec_f32 = false); + bool force_prec_f32 = false, + bool direct = false); // w: [OC，IC, KD, 1 * 1] // x: [N, IC, ID, IH*IW] diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp index 65ed8b87..bb08bb14 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -530,6 +530,7 @@ GGMLRunnerContext GGMLRunner::get_context() { runner_ctx.linear_scale = linear_scale; runner_ctx.attn_scale = attn_scale; runner_ctx.conv2d_direct_enabled = conv2d_direct_enabled; + runner_ctx.conv3d_direct_enabled = conv3d_direct_enabled; runner_ctx.circular_x_enabled = circular_x_enabled; runner_ctx.circular_y_enabled = circular_y_enabled; runner_ctx.weight_adapter = weight_adapter; diff --git a/src/core/ggml_runner.h b/src/core/ggml_runner.h index 5dc3fa5a..273098e8 100644 --- a/src/core/ggml_runner.h +++ b/src/core/ggml_runner.h @@ -71,6 +71,7 @@ struct GGMLRunnerContext { float linear_scale = 0.f; float attn_scale = 0.f; bool conv2d_direct_enabled = false; + bool conv3d_direct_enabled = false; bool circular_x_enabled = false; bool circular_y_enabled = false; ggml_tensor* ip_context = nullptr; @@ -178,6 +179,7 @@ protected: float linear_scale = 0.f; float attn_scale = 0.f; bool conv2d_direct_enabled = false; + bool conv3d_direct_enabled = false; bool circular_x_enabled = false; bool circular_y_enabled = false; @@ -346,6 +348,10 @@ public: conv2d_direct_enabled = enabled; } + void set_conv3d_direct_enabled(bool enabled) { + conv3d_direct_enabled = enabled; + } + void set_circular_axes(bool circular_x, bool circular_y) { circular_x_enabled = circular_x; circular_y_enabled = circular_y; diff --git a/src/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index 2aca72f4..54fa9d96 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -728,7 +728,7 @@ public: std::get<2>(stride), std::get<1>(stride), std::get<0>(stride), std::get<2>(padding), std::get<1>(padding), std::get<0>(padding), std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation), - force_prec_f32); + force_prec_f32, ctx->conv3d_direct_enabled); } }; diff --git a/src/model/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp index a56392e8..0ea8cd8c 100644 --- a/src/model/vae/wan_vae.hpp +++ b/src/model/vae/wan_vae.hpp @@ -78,7 +78,8 @@ namespace WAN { return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels, std::get<2>(stride), std::get<1>(stride), std::get<0>(stride), 0, 0, 0, - std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation)); + std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation), + false, ctx->conv3d_direct_enabled); } }; diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp index 40c60bf2..84d03278 100644 --- a/src/pipeline/model_builders.cpp +++ b/src/pipeline/model_builders.cpp @@ -584,10 +584,12 @@ namespace sd::model_builders { } if (sd_ctx_params->vae_conv_direct) { - LOG_INFO("Using Conv2d direct in the vae model"); + LOG_INFO("Using Conv2d/Conv3d direct in the vae model"); result.vae->set_conv2d_direct_enabled(true); + result.vae->set_conv3d_direct_enabled(true); if (result.preview) { result.preview->set_conv2d_direct_enabled(true); + result.preview->set_conv3d_direct_enabled(true); } } if (result.vae) {