Skip to content

Commit 94cc908

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 7f0e728 commit 94cc908

7 files changed

Lines changed: 144 additions & 9 deletions

File tree

examples/common/common.cpp

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -908,7 +908,7 @@ ArgOptions SDGenerationParams::get_options() {
908908
&extra_sample_args},
909909
{"",
910910
"--extra-tiling-args",
911-
"extra VAE tiling args, key=value list. LTX video VAE supports temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)",
911+
"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)",
912912
&extra_tiling_args},
913913
};
914914

@@ -1097,6 +1097,12 @@ ArgOptions SDGenerationParams::get_options() {
10971097
"process vae in tiles to reduce memory usage",
10981098
true,
10991099
&vae_tiling_params.enabled},
1100+
{"",
1101+
"--no-vae-tiling-fallback",
1102+
"disable the automatic fallback to VAE tiling when an untiled decode would exceed the "
1103+
"backend's max buffer size (fail instead of tiling)",
1104+
false,
1105+
&vae_tiling_params.auto_tile},
11001106
{"",
11011107
"--temporal-tiling",
11021108
"enable temporal tiling for LTX video VAE decode",
@@ -1841,6 +1847,9 @@ bool SDGenerationParams::from_json_str(
18411847
if (tiling_json.contains("enabled") && tiling_json["enabled"].is_boolean()) {
18421848
vae_tiling_params.enabled = tiling_json["enabled"];
18431849
}
1850+
if (tiling_json.contains("auto_tile") && tiling_json["auto_tile"].is_boolean()) {
1851+
vae_tiling_params.auto_tile = tiling_json["auto_tile"];
1852+
}
18441853
if (tiling_json.contains("temporal_tiling") && tiling_json["temporal_tiling"].is_boolean()) {
18451854
vae_tiling_params.temporal_tiling = tiling_json["temporal_tiling"];
18461855
}
@@ -2660,10 +2669,12 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
26602669
}
26612670

26622671
if (gen_params.vae_tiling_params.enabled ||
2672+
!gen_params.vae_tiling_params.auto_tile ||
26632673
gen_params.vae_tiling_params.temporal_tiling ||
26642674
!gen_params.extra_tiling_args.empty()) {
26652675
root["vae_tiling"] = {
26662676
{"enabled", gen_params.vae_tiling_params.enabled},
2677+
{"auto_tile", gen_params.vae_tiling_params.auto_tile},
26672678
{"temporal_tiling", gen_params.vae_tiling_params.temporal_tiling},
26682679
{"tile_size_x", gen_params.vae_tiling_params.tile_size_x},
26692680
{"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
@@ -228,7 +228,7 @@ struct SDGenerationParams {
228228
int video_frames = 1;
229229
int fps = 16;
230230
float vace_strength = 1.f;
231-
sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
231+
sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr, true};
232232
std::string extra_tiling_args;
233233

234234
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: 80 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,74 @@ 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+
// Defer a too-large untiled decode to tiling before the reserve hits a raw backend error.
1996+
if (probe_max_bytes_ > 0) {
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+
if (sizes[0] > probe_max_bytes_) {
2002+
LOG_DEBUG("%s: untiled compute buffer %.2f MB exceeds requested max_buffer_size %.2f MB; deferring to tiling",
2003+
get_desc().c_str(),
2004+
sizes[0] / 1024.0 / 1024.0,
2005+
probe_max_bytes_ / 1024.0 / 1024.0);
2006+
compute_buffer_deferred_to_tiling = true;
2007+
return false;
2008+
}
2009+
}
2010+
if (sd_backend_is(runtime_backend, "Vulkan")) {
2011+
// buft_get_max_size only reports Vulkan's suballocation block; supports_op has the real per-buffer limit.
2012+
for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) {
2013+
ggml_tensor* op = ggml_graph_node(gf, i);
2014+
if (!ggml_backend_supports_op(runtime_backend, op)) {
2015+
LOG_DEBUG("%s: untiled compute op %.2f MB exceeds backend support; deferring to tiling",
2016+
get_desc().c_str(),
2017+
ggml_nbytes(op) / 1024.0 / 1024.0);
2018+
compute_buffer_deferred_to_tiling = true;
2019+
return false;
2020+
}
2021+
}
2022+
} else {
2023+
ggml_gallocr* probe = ggml_gallocr_new(buft);
2024+
size_t sizes[1] = {0};
2025+
ggml_gallocr_reserve_n_size(probe, gf, nullptr, nullptr, sizes);
2026+
ggml_gallocr_free(probe);
2027+
size_t planned = sizes[0];
2028+
2029+
size_t max_size = ggml_backend_buft_get_max_size(buft);
2030+
bool over_buffer_cap = max_size > 0 && planned > max_size;
2031+
2032+
// CUDA/ROCm have no per-buffer cap, so gate on free VRAM plus a margin for the scratch pool the reserve omits.
2033+
bool over_free_vram = false;
2034+
ggml_backend_dev_t dev = ggml_backend_get_device(runtime_backend);
2035+
if (dev != nullptr && ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) {
2036+
size_t free_vram = 0, total_vram = 0;
2037+
ggml_backend_dev_memory(dev, &free_vram, &total_vram);
2038+
size_t margin = planned / 3;
2039+
if (margin < 512ull * 1024 * 1024) {
2040+
margin = 512ull * 1024 * 1024;
2041+
}
2042+
over_free_vram = free_vram > 0 && free_vram < planned + margin;
2043+
}
2044+
2045+
if (over_buffer_cap || over_free_vram) {
2046+
LOG_DEBUG("%s: untiled compute buffer %.2f MB won't fit free VRAM; deferring to tiling",
2047+
get_desc().c_str(),
2048+
planned / 1024.0 / 1024.0);
2049+
compute_buffer_deferred_to_tiling = true;
2050+
return false;
2051+
}
2052+
}
2053+
}
2054+
2055+
compute_allocr = ggml_gallocr_new(buft);
19852056

19862057
if (!ggml_gallocr_reserve(compute_allocr, gf)) {
19872058
// failed to allocate the compute buffer
@@ -2432,7 +2503,9 @@ struct GGMLRunner {
24322503
GraphWeightDoneGuard graph_weight_done_guard(this, &params_to_prepare);
24332504

24342505
if (!alloc_compute_buffer(gf)) {
2435-
LOG_ERROR("%s alloc compute buffer failed", get_desc().c_str());
2506+
if (!compute_buffer_deferred_to_tiling) {
2507+
LOG_ERROR("%s alloc compute buffer failed", get_desc().c_str());
2508+
}
24362509
return std::nullopt;
24372510
}
24382511
struct ComputeBufferGuard {
@@ -2822,6 +2895,11 @@ struct GGMLRunner {
28222895
void set_stream_layers_enabled(bool enabled) {
28232896
stream_layers_enabled = enabled;
28242897
}
2898+
2899+
void set_probe_compute_buffer_fits(bool enabled, size_t max_bytes = 0) {
2900+
probe_compute_buffer_fits_ = enabled;
2901+
probe_max_bytes_ = enabled ? max_bytes : 0;
2902+
}
28252903
};
28262904

28272905
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)