Skip to content

Commit 927734e

Browse files
author
Mark Caldwell
committed
feat: automatic VAE-tiling fallback when an untiled decode exceeds the backend buffer limit
A full-image VAE decode can need a single compute buffer bigger than the backend allows, so sd.cpp hard-failed instead of using the tiling it already has. This makes that fallback automatic and on by default, so a run no longer dies one flag short of working. --vae-tiling still forces tiling on; --no-vae-tiling-fallback restores the old hard-fail. Before allocating, the planned buffer is measured from the graph and checked per backend: Vulkan uses ggml_backend_supports_op (its real per-buffer limit), CUDA/ROCm check free VRAM since they report no per-buffer cap, keeping a margin for the compute pool the reserve doesn't count. If the untiled decode still comes back empty it frees and retries tiled, covering a genuine OOM. max_buffer_size in extra_tiling_args caps it manually on any backend. auto_tile is appended to sd_tiling_params_t so the C ABI stays compatible.
1 parent 2bd249c commit 927734e

7 files changed

Lines changed: 147 additions & 9 deletions

File tree

examples/common/common.cpp

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -958,7 +958,7 @@ ArgOptions SDGenerationParams::get_options() {
958958
&extra_sample_args},
959959
{"",
960960
"--extra-tiling-args",
961-
"extra VAE tiling args, key=value list. LTX video VAE supports temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)",
961+
"extra VAE tiling args, key=value list. max_buffer_size (bytes) forces the auto fallback to tile when an untiled VAE compute buffer would exceed it. LTX video VAE supports temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)",
962962
(int)',',
963963
&extra_tiling_args},
964964
};
@@ -1148,6 +1148,12 @@ ArgOptions SDGenerationParams::get_options() {
11481148
"process vae in tiles to reduce memory usage",
11491149
true,
11501150
&vae_tiling_params.enabled},
1151+
{"",
1152+
"--no-vae-tiling-fallback",
1153+
"disable the automatic fallback to VAE tiling when an untiled decode would exceed the "
1154+
"backend's max buffer size (fail instead of tiling)",
1155+
false,
1156+
&vae_tiling_params.auto_tile},
11511157
{"",
11521158
"--temporal-tiling",
11531159
"enable temporal tiling for LTX video VAE decode",
@@ -1892,6 +1898,9 @@ bool SDGenerationParams::from_json_str(
18921898
if (tiling_json.contains("enabled") && tiling_json["enabled"].is_boolean()) {
18931899
vae_tiling_params.enabled = tiling_json["enabled"];
18941900
}
1901+
if (tiling_json.contains("auto_tile") && tiling_json["auto_tile"].is_boolean()) {
1902+
vae_tiling_params.auto_tile = tiling_json["auto_tile"];
1903+
}
18951904
if (tiling_json.contains("temporal_tiling") && tiling_json["temporal_tiling"].is_boolean()) {
18961905
vae_tiling_params.temporal_tiling = tiling_json["temporal_tiling"];
18971906
}
@@ -2711,10 +2720,12 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
27112720
}
27122721

27132722
if (gen_params.vae_tiling_params.enabled ||
2723+
!gen_params.vae_tiling_params.auto_tile ||
27142724
gen_params.vae_tiling_params.temporal_tiling ||
27152725
!gen_params.extra_tiling_args.empty()) {
27162726
root["vae_tiling"] = {
27172727
{"enabled", gen_params.vae_tiling_params.enabled},
2728+
{"auto_tile", gen_params.vae_tiling_params.auto_tile},
27182729
{"temporal_tiling", gen_params.vae_tiling_params.temporal_tiling},
27192730
{"tile_size_x", gen_params.vae_tiling_params.tile_size_x},
27202731
{"tile_size_y", gen_params.vae_tiling_params.tile_size_y},

examples/common/common.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -229,7 +229,7 @@ struct SDGenerationParams {
229229
int video_frames = 1;
230230
int fps = 16;
231231
float vace_strength = 1.f;
232-
sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
232+
sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr, true};
233233
std::string extra_tiling_args;
234234

235235
std::string pm_id_images_dir;

examples/server/api.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -518,7 +518,7 @@ Shared default fields used by both `img_gen` and `vid_gen`:
518518
| `output_format` | `string` |
519519
| `output_compression` | `integer` |
520520

521-
`vae_tiling_params.extra_tiling_args` accepts a key=value list. For LTX video VAE temporal tiling, `temporal_tile_frames` defaults to `4` and `temporal_tile_overlap` defaults to `1`.
521+
`vae_tiling_params.extra_tiling_args` accepts a key=value list. `max_buffer_size` (bytes) forces the automatic tiling fallback when an untiled VAE compute buffer would exceed it. For LTX video VAE temporal tiling, `temporal_tile_frames` defaults to `4` and `temporal_tile_overlap` defaults to `1`.
522522

523523
`img_gen`-specific default fields:
524524

include/stable-diffusion.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -153,14 +153,15 @@ enum lora_apply_mode_t {
153153
};
154154

155155
typedef struct {
156-
bool enabled;
156+
bool enabled; // true => always tile (ON)
157157
bool temporal_tiling;
158158
int tile_size_x;
159159
int tile_size_y;
160160
float target_overlap;
161161
float rel_size_x;
162162
float rel_size_y;
163163
const char* extra_tiling_args;
164+
bool auto_tile; // AUTO (default): tile only when an untiled VAE decode would exceed the backend's max buffer size
164165
} sd_tiling_params_t;
165166

166167
typedef struct {

src/core/ggml_extend.hpp

Lines changed: 83 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1705,11 +1705,18 @@ struct GGMLRunner {
17051705

17061706
ggml_context* compute_ctx = nullptr;
17071707
ggml_gallocr* compute_allocr = nullptr;
1708+
// Set when alloc_compute_buffer deferred to tiling on purpose (not a failure).
1709+
bool compute_buffer_deferred_to_tiling = false;
17081710

17091711
size_t max_graph_vram_bytes = 0;
17101712
bool stream_layers_enabled = false;
17111713
size_t observed_max_effective_budget_ = 0;
17121714

1715+
// When set, alloc_compute_buffer declines a too-large untiled decode so VAE AUTO can tile.
1716+
bool probe_compute_buffer_fits_ = false;
1717+
// Optional user cap (bytes) to force tiling; 0 = no cap.
1718+
size_t probe_max_bytes_ = 0;
1719+
17131720
std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
17141721
std::weak_ptr<RunnerWeightManager> weight_manager;
17151722
std::unordered_set<const ggml_tensor*> kept_compute_param_tensor_set;
@@ -1978,10 +1985,77 @@ struct GGMLRunner {
19781985
}
19791986

19801987
bool alloc_compute_buffer(ggml_cgraph* gf) {
1988+
compute_buffer_deferred_to_tiling = false;
19811989
if (compute_allocr != nullptr) {
19821990
return true;
19831991
}
1984-
compute_allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_backend));
1992+
ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(runtime_backend);
1993+
1994+
if (probe_compute_buffer_fits_) {
1995+
// Measure the planned untiled compute buffer once (no allocation), then defer to
1996+
// tiling before the real reserve hits a raw backend error.
1997+
ggml_gallocr* probe = ggml_gallocr_new(buft);
1998+
size_t sizes[1] = {0};
1999+
ggml_gallocr_reserve_n_size(probe, gf, nullptr, nullptr, sizes);
2000+
ggml_gallocr_free(probe);
2001+
size_t planned = sizes[0];
2002+
2003+
// User cap (extra_tiling_args max_buffer_size), any backend.
2004+
if (probe_max_bytes_ > 0 && planned > probe_max_bytes_) {
2005+
LOG_DEBUG("%s: untiled compute buffer %.2f MB exceeds requested max_buffer_size %.2f MB; deferring to tiling",
2006+
get_desc().c_str(),
2007+
planned / 1024.0 / 1024.0,
2008+
probe_max_bytes_ / 1024.0 / 1024.0);
2009+
compute_buffer_deferred_to_tiling = true;
2010+
return false;
2011+
}
2012+
2013+
// Free VRAM, any non-CPU backend: a decode can fit every op's per-buffer cap yet
2014+
// still exceed total free VRAM. Margin covers the scratch pool the reserve omits.
2015+
ggml_backend_dev_t dev = ggml_backend_get_device(runtime_backend);
2016+
if (dev != nullptr && ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) {
2017+
size_t free_vram = 0, total_vram = 0;
2018+
ggml_backend_dev_memory(dev, &free_vram, &total_vram);
2019+
size_t margin = planned / 3;
2020+
if (margin < 512ull * 1024 * 1024) {
2021+
margin = 512ull * 1024 * 1024;
2022+
}
2023+
if (free_vram > 0 && free_vram < planned + margin) {
2024+
LOG_DEBUG("%s: untiled compute buffer %.2f MB won't fit free VRAM; deferring to tiling",
2025+
get_desc().c_str(),
2026+
planned / 1024.0 / 1024.0);
2027+
compute_buffer_deferred_to_tiling = true;
2028+
return false;
2029+
}
2030+
}
2031+
2032+
// Per-buffer cap: Vulkan via supports_op (the real limit; buft_get_max_size only
2033+
// reports the suballocation block there), other backends via buft_get_max_size.
2034+
if (sd_backend_is(runtime_backend, "Vulkan")) {
2035+
for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) {
2036+
ggml_tensor* op = ggml_graph_node(gf, i);
2037+
if (!ggml_backend_supports_op(runtime_backend, op)) {
2038+
LOG_DEBUG("%s: untiled compute op %.2f MB exceeds backend support; deferring to tiling",
2039+
get_desc().c_str(),
2040+
ggml_nbytes(op) / 1024.0 / 1024.0);
2041+
compute_buffer_deferred_to_tiling = true;
2042+
return false;
2043+
}
2044+
}
2045+
} else {
2046+
size_t max_size = ggml_backend_buft_get_max_size(buft);
2047+
if (max_size > 0 && planned > max_size) {
2048+
LOG_DEBUG("%s: untiled compute buffer %.2f MB exceeds backend max single buffer %.2f MB; deferring to tiling",
2049+
get_desc().c_str(),
2050+
planned / 1024.0 / 1024.0,
2051+
max_size / 1024.0 / 1024.0);
2052+
compute_buffer_deferred_to_tiling = true;
2053+
return false;
2054+
}
2055+
}
2056+
}
2057+
2058+
compute_allocr = ggml_gallocr_new(buft);
19852059

19862060
if (!ggml_gallocr_reserve(compute_allocr, gf)) {
19872061
// failed to allocate the compute buffer
@@ -2432,7 +2506,9 @@ struct GGMLRunner {
24322506
GraphWeightDoneGuard graph_weight_done_guard(this, &params_to_prepare);
24332507

24342508
if (!alloc_compute_buffer(gf)) {
2435-
LOG_ERROR("%s alloc compute buffer failed", get_desc().c_str());
2509+
if (!compute_buffer_deferred_to_tiling) {
2510+
LOG_ERROR("%s alloc compute buffer failed", get_desc().c_str());
2511+
}
24362512
return std::nullopt;
24372513
}
24382514
struct ComputeBufferGuard {
@@ -2822,6 +2898,11 @@ struct GGMLRunner {
28222898
void set_stream_layers_enabled(bool enabled) {
28232899
stream_layers_enabled = enabled;
28242900
}
2901+
2902+
void set_probe_compute_buffer_fits(bool enabled, size_t max_bytes = 0) {
2903+
probe_compute_buffer_fits_ = enabled;
2904+
probe_max_bytes_ = enabled ? max_bytes : 0;
2905+
}
28252906
};
28262907

28272908
class GGMLBlock {

src/model/vae/vae.hpp

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -199,7 +199,52 @@ struct VAE : public GGMLRunner {
199199
"vae decode compute failed while processing a tile",
200200
silent);
201201
} else {
202+
// AUTO: probe first so a too-large decode tiles instead of erroring; output.empty() backstops a real OOM.
203+
const bool auto_probe = !tiling_params.enabled && tiling_params.auto_tile;
204+
if (auto_probe) {
205+
size_t max_bytes = 0;
206+
if (tiling_params.extra_tiling_args != nullptr) {
207+
for (const auto& [key, value] : parse_key_value_args(tiling_params.extra_tiling_args, "VAE extra tiling arg")) {
208+
if (key == "max_buffer_size") {
209+
max_bytes = strtoull(value.c_str(), nullptr, 10);
210+
}
211+
}
212+
}
213+
set_probe_compute_buffer_fits(true, max_bytes);
214+
}
202215
output = _compute(n_threads, input, true);
216+
if (auto_probe) {
217+
set_probe_compute_buffer_fits(false);
218+
}
219+
if (output.empty() && !tiling_params.enabled && tiling_params.auto_tile) {
220+
free_compute_buffer();
221+
if (!silent) {
222+
LOG_WARN("vae: untiled decode buffer exceeded the backend limit; retrying with tiling");
223+
}
224+
sd_tiling_params_t auto_tiling = tiling_params;
225+
auto_tiling.enabled = true;
226+
set_tiling_params(auto_tiling);
227+
const int scale_factor = get_scale_factor();
228+
int64_t W = input.shape()[0] * scale_factor;
229+
int64_t H = input.shape()[1] * scale_factor;
230+
float tile_overlap;
231+
int tile_size_x, tile_size_y;
232+
get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, auto_tiling, input.shape()[0], input.shape()[1]);
233+
output = tiled_compute(
234+
input,
235+
n_threads,
236+
static_cast<int>(W),
237+
static_cast<int>(H),
238+
scale_factor,
239+
tile_size_x,
240+
tile_size_y,
241+
tile_overlap,
242+
circular_x,
243+
circular_y,
244+
true,
245+
"vae decode compute failed while processing a tile",
246+
silent);
247+
}
203248
}
204249

205250
free_compute_buffer();

src/stable-diffusion.cpp

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -192,7 +192,7 @@ class StableDiffusionGGML {
192192
bool apply_lora_immediately = false;
193193

194194
std::string taesd_path;
195-
sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr};
195+
sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr, true};
196196
bool enable_mmap = false;
197197
sd::ggml_graph_cut::MaxVramAssignment max_vram_assignment;
198198
bool stream_layers = false;
@@ -2843,7 +2843,7 @@ void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) {
28432843
sd_img_gen_params->control_strength = 0.9f;
28442844
sd_img_gen_params->pm_params = {nullptr, 0, nullptr, 20.f};
28452845
sd_img_gen_params->pulid_params = {nullptr, 1.0f};
2846-
sd_img_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
2846+
sd_img_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr, true};
28472847
sd_cache_params_init(&sd_img_gen_params->cache);
28482848
sd_hires_params_init(&sd_img_gen_params->hires);
28492849
}
@@ -2930,7 +2930,7 @@ void sd_vid_gen_params_init(sd_vid_gen_params_t* sd_vid_gen_params) {
29302930
sd_vid_gen_params->fps = 16;
29312931
sd_vid_gen_params->moe_boundary = 0.875f;
29322932
sd_vid_gen_params->vace_strength = 1.f;
2933-
sd_vid_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
2933+
sd_vid_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr, true};
29342934
sd_vid_gen_params->hires.enabled = false;
29352935
sd_vid_gen_params->hires.upscaler = SD_HIRES_UPSCALER_LATENT;
29362936
sd_vid_gen_params->hires.scale = 2.f;

0 commit comments

Comments
 (0)