perf: accelerate VAE direct 3D convolutions (#1996)

This commit is contained in:
leejet 2026-09-19 17:50:01 +08:00 committed by GitHub
parent 2ea8aff7ef
commit 3e037a81e4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 24 additions and 8 deletions

View File

@ -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},
};

2
ggml

@ -1 +1 @@
Subproject commit 1e22ec0d04b43afa69963e4b4ea6683535d54d79
Subproject commit c6632cd905401abc58b6f5cdd52d228aa7ca1b88

View File

@ -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;

View File

@ -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]

View File

@ -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;

View File

@ -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;

View File

@ -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);
}
};

View File

@ -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);
}
};

View File

@ -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) {