From 58d3f6e240a7e5886e5829145311799dc1610fea Mon Sep 17 00:00:00 2001 From: leejet Date: Fri, 11 Sep 2026 01:05:30 +0800 Subject: [PATCH] fix: guard GPU memory capacity and propagate encoding failures --- docs/backend.md | 5 +++ src/conditioning/conditioner.hpp | 10 ++++-- src/core/ggml_runner.cpp | 12 +++++-- src/model_manager.cpp | 55 +++++++++++++++++++++----------- src/pipeline/image.cpp | 12 +++++++ 5 files changed, 72 insertions(+), 22 deletions(-) diff --git a/docs/backend.md b/docs/backend.md index edf7784da..429e659e7 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -149,6 +149,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: diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index 23cfbddd2..f0a7286f9 100644 --- a/src/conditioning/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -1224,7 +1224,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::zeros({768}); } @@ -1243,7 +1246,10 @@ struct FluxCLIPEmbedder : public Conditioner { input_ids, sd::Tensor(), 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); diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp index 3737dc7aa..22afc47f9 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -1,4 +1,5 @@ #include +#include #include #include @@ -624,8 +625,15 @@ std::optional> 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> 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); } diff --git a/src/model_manager.cpp b/src/model_manager.cpp index f3bf19a2c..734eb053e 100644 --- a/src/model_manager.cpp +++ b/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; } diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp index 9591c5d60..69c487421 100644 --- a/src/pipeline/image.cpp +++ b/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 } @@ -466,6 +470,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 @@ -491,6 +499,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 }