fix: guard GPU memory capacity and propagate encoding failures (#1958)

CMakeLists.txt

@@ -294,6 +294,8 @@ endif()
 
 if(MSVC)
     target_compile_options(${SD_LIB} PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/bigobj>)
+    # ggml backends can throw C++ exceptions through their C API.
+    target_compile_options(${SD_LIB} PRIVATE $<$<AND:$<COMPILE_LANGUAGE:CXX>,$<CXX_COMPILER_ID:MSVC>>:/EHsc->)
 endif()
 
 if(APPLE)

docs/backend.md

@@ -154,6 +154,11 @@ GiB", and with no budget set each device's free memory minus a 512 MiB margin
 is used. These resolved GPU budgets, including the safety margin, also drive
 the runner's graph-cut capacity checks.
 
+Runtime capacity checks also leave 512 MiB of currently free device memory for
+backend scratch buffers and pipelines, including with explicit backend assignments.
+They cap stale free-memory reports by the device's total memory minus tracked
+resident allocations and reject reports that exceed the device's total memory.
+
 Components are considered in `diffusion`, `te`, `vae` order so that repeatedly
 used diffusion weights have priority. Each component's weights use the first
 storage location with enough remaining budget:

src/conditioning/conditioner.hpp

@@ -1254,7 +1254,10 @@ struct FluxCLIPEmbedder : public Conditioner {
                                              true,
                                              clip_skip,
                                              false);
-                    GGML_ASSERT(!pooled.empty());
+                    if (pooled.empty()) {
+                        LOG_ERROR("Flux CLIP-L encoding failed");
+                        return {};
+                    }
                 } else {
                     pooled = sd::Tensor<float>::zeros({768});
                 }
@@ -1273,7 +1276,10 @@ struct FluxCLIPEmbedder : public Conditioner {
                                                   input_ids,
                                                   sd::Tensor<float>(),
                                                   false);
-                GGML_ASSERT(!chunk_hidden_states.empty());
+                if (chunk_hidden_states.empty()) {
+                    LOG_ERROR("Flux T5 encoding failed at chunk %d/%zu", chunk_idx + 1, chunk_count);
+                    return {};
+                }
                 chunk_hidden_states = ::apply_token_weights(std::move(chunk_hidden_states), chunk_weights);
                 if (zero_out_masked) {
                     chunk_hidden_states.fill_(0.0f);

src/core/compute_workspace.cpp

@@ -2,12 +2,14 @@
 
 #include <algorithm>
 #include <cstring>
+#include <exception>
 #include <map>
 #include <unordered_map>
 #include <unordered_set>
 
 #include "core/ggml_extend_backend.h"
 #include "core/ggml_graph_cut.h"
+#include "core/util.h"
 #include "ggml-cpu.h"
 #include "ggml/src/ggml-impl.h"
 
@@ -228,11 +230,23 @@ namespace sd {
         }
     }
 
-    void ComputeWorkspace::segment_end() {
-        if (active_) {
-            synchronize();
-            active_ = false;
+    bool ComputeWorkspace::segment_end() noexcept {
+        if (!active_) {
+            return true;
         }
+        // Outer cleanup guards must not retry a failed backend submission.
+        active_ = false;
+        try {
+            synchronize();
+            return true;
+        } catch (const std::exception& error) {
+            LOG_ERROR("%s workspace synchronization failed during segment cleanup: %s",
+                      ggml_backend_name(backend_), error.what());
+        } catch (...) {
+            LOG_ERROR("%s workspace synchronization failed during segment cleanup: unknown exception",
+                      ggml_backend_name(backend_));
+        }
+        return false;
     }
 
     bool ComputeWorkspace::release() {

src/core/compute_workspace.h

@@ -51,7 +51,7 @@ namespace sd {
             const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
             const AssignNodes& assign_nodes);
         void synchronize() const;
-        void segment_end();
+        bool segment_end() noexcept;
         bool release();
         bool active() const { return active_; }
         ggml_backend_sched_t scheduler() const { return scheduler_; }

src/core/ggml_extend_backend.cpp

@@ -660,7 +660,13 @@ void SDBackendAssignment::set_module(SDBackendModule module, const std::string&
 }
 
 void SDBackendHandleDeleter::operator()(ggml_backend_t backend) const {
-    ggml_backend_free(backend);
+    try {
+        ggml_backend_free(backend);
+    } catch (const std::exception& error) {
+        LOG_ERROR("backend cleanup failed: %s", error.what());
+    } catch (...) {
+        LOG_ERROR("backend cleanup failed: unknown exception");
+    }
 }
 
 SDBackendManager::~SDBackendManager() {

src/core/ggml_runner.cpp

@@ -1,4 +1,5 @@
 #include <algorithm>
+#include <exception>
 #include <map>
 #include <utility>
 
@@ -642,8 +643,15 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
                 params_tensor_set_.insert(parameter);
         }
     }
-    auto output = execute_graph(graph, n_threads, no_return, read_outputs);
-    success     = output.has_value();
+    std::optional<sd::Tensor<float>> output;
+    try {
+        output = execute_graph(graph, n_threads, no_return, read_outputs);
+    } catch (const std::exception& error) {
+        LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(),
+                  ggml_backend_name(runtime_backend), error.what());
+        return std::nullopt;
+    }
+    success = output.has_value();
     if (success) {
         cache_.graph_end(true);
     }
@@ -956,6 +964,9 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
                 }
             }
         }
+        if (!workspace_.segment_end()) {
+            return fail_segment("workspace synchronization");
+        }
         // Final outputs and their callbacks may still be views of consumed cuts.
         cut_cache_.prune(segment.future_cut_names);
     }

src/model_manager.cpp

@@ -1584,18 +1584,35 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
     if (request.compute_backend == nullptr || sd_backend_is_cpu(request.compute_backend)) {
         return result;
     }
-    auto add                     = [](size_t a, size_t b) { return b > SIZE_MAX - a ? SIZE_MAX : a + b; };
-    const size_t missing         = compute_backend_alloc_size(states, true);
-    result.required_device_bytes = add(request.pending_allocation_bytes, missing);
-    result.required_budget_bytes = add(request.runtime_peak_bytes(), missing);
-    auto device                  = ggml_backend_get_device(request.compute_backend);
-    if (device != nullptr) {
+    auto add             = [](size_t a, size_t b) { return b > SIZE_MAX - a ? SIZE_MAX : a + b; };
+    const size_t missing = compute_backend_alloc_size(states, true);
+    // Backend scratch buffers and pipelines are not included in graph measurements.
+    constexpr size_t safety_margin = 512ULL * 1024ULL * 1024ULL;
+    result.required_device_bytes   = add(add(request.pending_allocation_bytes, missing), safety_margin);
+    result.required_budget_bytes   = add(request.runtime_peak_bytes(), missing);
+    auto available_device_bytes    = [&](ggml_backend_t backend) {
+        auto device = ggml_backend_get_device(backend);
+        if (device == nullptr) {
+            return SIZE_MAX;
+        }
         size_t free_bytes = 0, total_bytes = 0;
         ggml_backend_dev_memory(device, &free_bytes, &total_bytes);
-        if (free_bytes != 0 || total_bytes != 0) {
-            result.available_device_bytes = free_bytes;
+        if (free_bytes == 0 && total_bytes == 0) {
+            return SIZE_MAX;
         }
-    }
+        // Vulkan's heap budget subtraction can underflow when usage exceeds the budget.
+        if (total_bytes > 0 && free_bytes > total_bytes) {
+            return size_t{0};
+        }
+        const size_t resident = add(compute_backend_resident_bytes(backend),
+                                       add(other_runtime_resident_bytes(request.owner_id, backend),
+                                           request.runtime_resident_bytes));
+        if (total_bytes > 0) {
+            free_bytes = std::min(free_bytes, resident < total_bytes ? total_bytes - resident : 0);
+        }
+        return free_bytes;
+    };
+    result.available_device_bytes = available_device_bytes(request.compute_backend);
     if (request.max_backend_bytes > 0) {
         const size_t resident         = add(compute_backend_resident_bytes(request.compute_backend),
                                             other_runtime_resident_bytes(request.owner_id, request.compute_backend));
@@ -1619,11 +1636,7 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
     // GGML exposes only a split buffer's total size, not per-device allocations.
     // Charge that upper bound on every participant instead of undercounting a shard.
     for (const auto& entry : split_devices) {
-        size_t free_bytes = 0, total_bytes = 0;
-        ggml_backend_dev_memory(ggml_backend_get_device(entry.first), &free_bytes, &total_bytes);
-        if (free_bytes != 0 || total_bytes != 0) {
-            result.available_device_bytes = std::min(result.available_device_bytes, free_bytes);
-        }
+        result.available_device_bytes = std::min(result.available_device_bytes, available_device_bytes(entry.first));
         if (entry.second > 0) {
             const size_t resident         = add(compute_backend_resident_bytes(entry.first),
                                                 other_runtime_resident_bytes(request.owner_id, entry.first));
@@ -1739,12 +1752,18 @@ bool ModelManager::ensure_compute_backend_capacity(
         }
     }
 
-    const auto capacity = check_capacity(request, required_states);
-    LOG_WARN("model manager cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %.2f MB device / %.2f MB budget",
+    const auto capacity                = check_capacity(request, required_states);
+    const std::string available_device = capacity.available_device_bytes == SIZE_MAX
+                                             ? "unknown"
+                                             : sd_format("%.2f MB", capacity.available_device_bytes / (1024.0 * 1024.0));
+    const std::string available_budget = capacity.available_budget_bytes == SIZE_MAX
+                                             ? "unlimited"
+                                             : sd_format("%.2f MB", capacity.available_budget_bytes / (1024.0 * 1024.0));
+    LOG_WARN("model manager cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %s device / %s budget",
              ggml_backend_name(compute_backend),
              capacity.required_device_bytes / (1024.0 * 1024.0),
              capacity.required_budget_bytes / (1024.0 * 1024.0),
-             capacity.available_device_bytes / (1024.0 * 1024.0),
-             capacity.available_budget_bytes / (1024.0 * 1024.0));
+             available_device.c_str(),
+             available_budget.c_str());
     return false;
 }

src/model_manager.h

@@ -31,7 +31,7 @@ public:
     };
 
 private:
-    static constexpr size_t MAX_RESIDENCY_BLOCK_BYTES = 64ULL * 1024ULL * 1024ULL;
+    static constexpr size_t MAX_RESIDENCY_BLOCK_BYTES = 1024ULL * 1024ULL * 1024ULL;
 
     struct TensorState {
         std::string name;

src/pipeline/image.cpp

@@ -438,6 +438,10 @@ namespace sd::pipeline {
         condition_params.zero_out_masked = false;
         auto cond                        = sd->cond_stage_model->get_learned_condition(sd->n_threads,
                                                                                        condition_params);
+        if (cond.empty()) {
+            LOG_ERROR("failed to encode prompt");
+            return std::nullopt;
+        }
         if (cond.c_concat.empty() && ref_image_params.pass_to_dit) {
             cond.c_concat = latents->concat_latent;  // TODO: optimize
         }
@@ -469,6 +473,10 @@ namespace sd::pipeline {
                 condition_params.zero_out_masked = zero_out_masked;
                 uncond                           = sd->cond_stage_model->get_learned_condition(sd->n_threads,
                                                                                                condition_params);
+                if (uncond.empty()) {
+                    LOG_ERROR("failed to encode negative prompt");
+                    return std::nullopt;
+                }
             }
             if (uncond.c_concat.empty() && ref_image_params.pass_to_dit) {
                 uncond.c_concat = latents->concat_latent;  // TODO: optimize
@@ -494,6 +502,10 @@ namespace sd::pipeline {
                 }
                 img_uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
                                                                          condition_params);
+                if (img_uncond.empty()) {
+                    LOG_ERROR("failed to encode image guidance prompt");
+                    return std::nullopt;
+                }
                 if (img_uncond.c_concat.empty() && ref_image_params.pass_to_dit) {
                     img_uncond.c_concat = latents->img_uncond_concat_latent;  // TODO: optimize
                 }