From 97d932b8f81264165cee0285fd62484ce67952d6 Mon Sep 17 00:00:00 2001 From: leejet Date: Mon, 21 Sep 2026 23:40:05 +0800 Subject: [PATCH] fix: restrict VAE tiling retries to allocation failures (#2019) --- examples/common/common.cpp | 2 +- src/core/backend_fit.cpp | 8 ++++++-- src/core/backend_fit.h | 3 ++- src/core/ggml_runner.cpp | 28 ++++++++++++++++++++++++---- src/core/ggml_runner.h | 5 ++++- src/pipeline/diffusion_engine.cpp | 3 ++- 6 files changed, 39 insertions(+), 10 deletions(-) diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 5a04d792..cc97f669 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -1754,7 +1754,7 @@ ArgOptions SDGenerationParams::get_options() { on_scm_policy_arg}, {"", "--vae-tile-size", - "tile size for vae tiling, format [X]x[Y] (default: 32x32)", + "tile size for vae tiling in latent units, not image pixels, format [X]x[Y] (default: 32x32)", on_tile_size_arg}, {"", "--vae-relative-tile-size", diff --git a/src/core/backend_fit.cpp b/src/core/backend_fit.cpp index 85cfc727..2ffc1cbb 100644 --- a/src/core/backend_fit.cpp +++ b/src/core/backend_fit.cpp @@ -478,7 +478,11 @@ namespace sd::backend_fit { return true; } - bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling) { + bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling, ggml_status status) { + // Execution failures can leave the device unusable; tiling only helps with allocation failures. + if (status != GGML_STATUS_ALLOC_FAILED) { + return false; + } const char* retry_mode = nullptr; if (prefer_temporal_tiling && !tiling_params.temporal_tiling) { tiling_params.temporal_tiling = true; @@ -498,7 +502,7 @@ namespace sd::backend_fit { return false; } - LOG_WARN("VAE decode failed (likely out of memory); retrying with %s tiling", + LOG_WARN("VAE decode ran out of memory; retrying with %s tiling", retry_mode); return true; } diff --git a/src/core/backend_fit.h b/src/core/backend_fit.h index 9ef298b3..0c682a40 100644 --- a/src/core/backend_fit.h +++ b/src/core/backend_fit.h @@ -16,7 +16,8 @@ namespace sd::backend_fit { std::string& params_spec); bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, - bool prefer_temporal_tiling); + bool prefer_temporal_tiling, + ggml_status status); } // namespace sd::backend_fit diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp index abfa8639..182fd098 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -590,6 +590,7 @@ std::optional> GGMLRunner::compute(get_graph_cb_t get_graph, bool auto_runner_end, bool no_return, const std::function& read_outputs) { + last_compute_status_ = GGML_STATUS_FAILED; if (graph_active_) { LOG_ERROR("%s does not support reentrant graph execution", get_desc().c_str()); return std::nullopt; @@ -613,7 +614,9 @@ std::optional> GGMLRunner::compute(get_graph_cb_t get_graph, GGMLRunner& runner; const bool& success; ~GraphEndGuard() { - runner.workspace_.segment_end(); + if (!runner.workspace_.segment_end()) { + runner.last_compute_status_ = GGML_STATUS_FAILED; + } runner.cache_.graph_end(false); runner.cut_cache_.clear(); runner.free_compute_ctx(); @@ -642,6 +645,7 @@ std::optional> GGMLRunner::compute(get_graph_cb_t get_graph, try { output = execute_graph(graph, n_threads, no_return, read_outputs); } catch (const std::exception& error) { + last_compute_status_ = GGML_STATUS_FAILED; LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(), ggml_backend_name(runtime_backend), error.what()); return std::nullopt; @@ -649,6 +653,7 @@ std::optional> GGMLRunner::compute(get_graph_cb_t get_graph, success = output.has_value(); if (success) { cache_.graph_end(true); + last_compute_status_ = GGML_STATUS_SUCCESS; } return output; } @@ -766,6 +771,7 @@ bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads) { } workspace_.synchronize(); if (status != GGML_STATUS_SUCCESS) { + last_compute_status_ = status; LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status)); return false; } @@ -818,6 +824,7 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n const auto& cached_plan = resolve_graph_cut_plan(graph); const auto full_measurement = measure(graph, cached_plan.compute_buffer_size); if (full_measurement.buffers.empty()) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; return std::nullopt; } auto manager = residency_manager.lock(); @@ -888,7 +895,9 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n SegmentGraphBindings& bindings; ggml_context* context; ~SegmentCleanup() { - runner.workspace_.segment_end(); + if (!runner.workspace_.segment_end()) { + runner.last_compute_status_ = GGML_STATUS_FAILED; + } bindings.restore(); weights.segment_end(); ggml_free(context); @@ -898,6 +907,7 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n auto measurement = segmented ? measure(segment_graph, segment.compute_buffer_size) : full_measurement; if (!workspace_.prepare(measurement)) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; return fail_segment("workspace preparation"); } const size_t cut_bytes = last ? 0 : cut_cache_.estimate_output_bytes(graph, segment); @@ -912,7 +922,11 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n sync_runtime_residency(); requests = memory_requests(measurement.buffers, new_cache_bytes); } - return weights.ensure_segment_capacity(index, requests); + const bool ready = weights.ensure_segment_capacity(index, requests); + if (!ready && manager != nullptr) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + } + return ready; }; if (!weights.segment_start(index, ensure_capacity)) { return fail_segment("weight preparation"); @@ -921,12 +935,17 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n if (!workspace_.measurement_matches(segment_graph, measurement)) { measurement = measure(segment_graph, segment.compute_buffer_size); } - if (!workspace_.prepare(measurement) || !ensure_capacity()) { + if (!workspace_.prepare(measurement)) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + return fail_segment("workspace preparation"); + } + if (!ensure_capacity()) { return fail_segment("workspace capacity check"); } if (!workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) { pin_multi_device_nodes(scheduler, current); })) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; return fail_segment("workspace allocation"); } for (const auto& size : measurement.buffers) { @@ -964,6 +983,7 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n } } if (!workspace_.segment_end()) { + last_compute_status_ = GGML_STATUS_FAILED; return fail_segment("workspace synchronization"); } // Final outputs and their callbacks may still be views of consumed cuts. diff --git a/src/core/ggml_runner.h b/src/core/ggml_runner.h index 29a6808b..9084db10 100644 --- a/src/core/ggml_runner.h +++ b/src/core/ggml_runner.h @@ -130,7 +130,8 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx, struct GGMLRunner { private: std::map logged_compute_bytes_; - size_t logged_segment_count_ = 0; + size_t logged_segment_count_ = 0; + ggml_status last_compute_status_ = GGML_STATUS_SUCCESS; sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, size_t direct_bytes); std::vector memory_requests(const std::vector& sizes, @@ -335,6 +336,8 @@ public: bool no_return = false, const std::function& read_outputs = {}); + ggml_status last_compute_status() const { return last_compute_status_; } + void set_flash_attention_enabled(bool enabled) { flash_attn_enabled = enabled; } diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index e41c1151..e020a2d6 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -2766,7 +2766,8 @@ sd::Tensor StableDiffusionGGML::decode_first_stage(const sd::Tensordecode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); const bool prefer_temporal_tiling = decode_video && first_stage_model->can_temporal_tile_decode(); while (decoded.empty() && - sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) { + sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling, + first_stage_model->last_compute_status())) { decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); } return decoded;