fix: make capacity checks account for host-pinned parameter memory (#1976)

src/core/ggml_runner.cpp

@@ -3,6 +3,12 @@
 #include <map>
 #include <utility>
 
+#ifdef _WIN32
+    #define NOMINMAX
+    #include <windows.h>
+    #include <psapi.h>
+#endif
+
 #include "core/ggml_extend.h"
 #include "core/ggml_extend_backend.h"
 #include "core/ggml_runner.h"
@@ -983,8 +989,60 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
         }
         LOG_DEBUG("%s executing segment %zu/%zu: %s", get_desc().c_str(),
                   index + 1, plan.segments.size(), segment.group_name.c_str());
-        if (!execute_segment(segment_graph, n_threads)) {
-            return fail_segment("execution");
+        {
+            size_t dev_free_bytes  = 0;
+            size_t dev_total_bytes = 0;
+            auto   dev             = ggml_backend_get_device(runtime_backend);
+            if (dev != nullptr) {
+                ggml_backend_dev_memory(dev, &dev_free_bytes, &dev_total_bytes);
+            }
+            size_t cpu_free_bytes  = 0;
+            size_t cpu_total_bytes = 0;
+            if (workspace_.cpu_backend() != nullptr) {
+                if (auto cpu_dev = ggml_backend_get_device(workspace_.cpu_backend()); cpu_dev != nullptr) {
+                    ggml_backend_dev_memory(cpu_dev, &cpu_free_bytes, &cpu_total_bytes);
+                }
+            }
+            const double cut_mb =
+                (dev != nullptr ? cut_cache_.resident_bytes(dev) : 0) / (1024.0 * 1024.0);
+#ifdef _WIN32
+            MEMORYSTATUSEX           mem_status{};
+            PROCESS_MEMORY_COUNTERS_EX proc_counters{};
+            mem_status.dwLength    = sizeof(mem_status);
+            proc_counters.cb       = sizeof(proc_counters);
+            const bool have_status = GlobalMemoryStatusEx(&mem_status) != 0;
+            const bool have_proc   = GetProcessMemoryInfo(GetCurrentProcess(),
+                                                reinterpret_cast<PROCESS_MEMORY_COUNTERS*>(&proc_counters),
+                                                sizeof(proc_counters)) != 0;
+            LOG_VERBOSE("%s segment %zu/%zu start on %s: device reports %.2f MB free / %.2f MB total"
+                        " | host: avail phys %.2f MB, avail commit %.2f MB, proc commit %.2f MB, proc ws %.2f MB"
+                        " | CPU device %.2f MB free / %.2f MB total | cut cache %.2f MB",
+                        get_desc().c_str(), index + 1, plan.segments.size(),
+                        ggml_backend_name(runtime_backend),
+                        dev_free_bytes / (1024.0 * 1024.0), dev_total_bytes / (1024.0 * 1024.0),
+                        have_status ? mem_status.ullAvailPhys / (1024.0 * 1024.0) : -1.0,
+                        have_status ? mem_status.ullAvailPageFile / (1024.0 * 1024.0) : -1.0,
+                        have_proc ? proc_counters.PrivateUsage / (1024.0 * 1024.0) : -1.0,
+                        have_proc ? proc_counters.WorkingSetSize / (1024.0 * 1024.0) : -1.0,
+                        cpu_free_bytes / (1024.0 * 1024.0), cpu_total_bytes / (1024.0 * 1024.0), cut_mb);
+#else
+            LOG_VERBOSE("%s segment %zu/%zu start on %s: device reports %.2f MB free / %.2f MB total"
+                        " | CPU device %.2f MB free / %.2f MB total | cut cache %.2f MB",
+                        get_desc().c_str(), index + 1, plan.segments.size(),
+                        ggml_backend_name(runtime_backend),
+                        dev_free_bytes / (1024.0 * 1024.0), dev_total_bytes / (1024.0 * 1024.0),
+                        cpu_free_bytes / (1024.0 * 1024.0), cpu_total_bytes / (1024.0 * 1024.0), cut_mb);
+#endif
+        }
+        try {
+            if (!execute_segment(segment_graph, n_threads)) {
+                return fail_segment("execution");
+            }
+        } catch (const std::exception& error) {
+            LOG_VERBOSE("%s segment %zu/%zu (%s) raised during execution: %s",
+                        get_desc().c_str(), index + 1, plan.segments.size(),
+                        segment.group_name.c_str(), error.what());
+            throw;
         }
         auto cache_status = cache_.capture(segment_graph);
         if (cache_status == GGML_STATUS_SUCCESS) {

src/core/segment_weight_pipeline.cpp

@@ -3,6 +3,7 @@
 #include <utility>
 
 #include "core/ggml_graph_cut.h"
+#include "core/util.h"
 #include "device_residency_manager.h"
 
 namespace sd {
@@ -67,15 +68,22 @@ namespace sd {
 
     void SegmentWeightPipeline::activate(size_t segment_index) {
         if (!enabled_ || queued_segment_ == SIZE_MAX || queued_segment_ != segment_index) {
+            if (enabled_ && queued_segment_ != SIZE_MAX) {
+                LOG_VERBOSE("segment weight pipeline: segment %zu started without its prefetched block (queued for segment %zu)",
+                            segment_index, queued_segment_);
+            }
             return;
         }
 
         auto manager = residency_manager_.lock();
         if (manager == nullptr ||
             !manager->activate_prefetched_params(owner_id_, queued_params_)) {
+            LOG_VERBOSE("segment weight pipeline: failed to activate the prefetched block for segment %zu, disabling prefetch",
+                        segment_index);
             disable();
             return;
         }
+        LOG_VERBOSE("segment weight pipeline: segment %zu activated its prefetched weight block", segment_index);
         queued_params_.clear();
         queued_segment_ = SIZE_MAX;
     }
@@ -175,6 +183,8 @@ namespace sd {
                                                       params,
                                                       preferred_eviction_order(),
                                                       protected_params)) {
+            LOG_VERBOSE("segment weight pipeline: prefetch capacity check failed for segment %zu, disabling prefetch",
+                        next_segment);
             disable();
             return;
         }

src/model_manager.cpp

@@ -7,6 +7,11 @@
 #include <tuple>
 #include <unordered_set>
 
+#ifdef _WIN32
+    #define NOMINMAX
+    #include <windows.h>
+#endif
+
 #include "core/ggml_extend_backend.h"
 #include "core/util.h"
 #include "model/adapter/lora.hpp"
@@ -17,6 +22,20 @@ static size_t aligned_offset(const void* buffer, size_t offset, size_t alignment
     return offset + align;
 }
 
+// True when allocations through this buffer type consume host-visible memory:
+// plain host buffer types, or a compute device's host-pinned pool (e.g. the
+// Vulkan "Vulkan_Host" buffer type backed by pinned system memory).
+static bool model_manager_buffer_type_is_host_pinned(ggml_backend_buffer_type_t buft) {
+    if (buft == nullptr) {
+        return false;
+    }
+    if (ggml_backend_buft_is_host(buft)) {
+        return true;
+    }
+    ggml_backend_dev_t device = ggml_backend_buft_get_device(buft);
+    return device != nullptr && buft == ggml_backend_dev_host_buffer_type(device);
+}
+
 static bool lora_specs_equal(const std::vector<ModelManager::LoraSpec>& lhs,
                              const std::vector<ModelManager::LoraSpec>& rhs) {
     if (lhs.size() != rhs.size()) {
@@ -1780,6 +1799,11 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
         const size_t weights_resident = compute_backend_resident_bytes(backend);
         const size_t other_runtime    = other_runtime_resident_bytes(request.owner_id, backend);
         const size_t resident         = add(weights_resident, add(other_runtime, request.runtime_resident_bytes));
+        LOG_VERBOSE("model manager memory on %s: reported free %.2f MB / total %.2f MB, tracked weights %.2f MB / other runtime %.2f MB / current runtime %.2f MB",
+                    ggml_backend_name(backend),
+                    free_bytes / (1024.0 * 1024.0), total_bytes / (1024.0 * 1024.0),
+                    weights_resident / (1024.0 * 1024.0), other_runtime / (1024.0 * 1024.0),
+                    request.runtime_resident_bytes / (1024.0 * 1024.0));
         if (log_details) {
             LOG_WARN("model manager memory on %s: reported free %.2f MB / total %.2f MB, tracked weights %.2f MB / other runtime %.2f MB / current runtime %.2f MB",
                         ggml_backend_name(backend),
@@ -1838,6 +1862,82 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
             result.available_budget_bytes = std::min(result.available_budget_bytes, available);
         }
     }
+    // Host dimension (#1976): memory allocated through host-visible buffer
+    // types is invisible to the compute device's heap report. With an
+    // offloaded model (params backend != compute backend), params buffers
+    // take the compute device's host-pinned buffer type (e.g. Vulkan_Host)
+    // and are loaded lazily, so they accumulate on the host while every
+    // device-side gate stays green. Account for the actual placement: only
+    // bytes this request would allocate through host-visible buffer types
+    // count against host availability; device-local staging/prefetch stay on
+    // the device/budget dimensions. A host failure feeds the existing
+    // eviction path, which reclaims params storage blocks and staged replicas.
+    size_t host_missing_bytes = 0;
+    bool   host_pool_in_use   = false;
+    // Existing host-pinned parameter storage keeps the floor active even when
+    // this request pins nothing new.
+    for (const auto& block : params_storage_blocks_) {
+        if (block != nullptr && block->buffer != nullptr &&
+            model_manager_buffer_type_is_host_pinned(ggml_backend_buffer_get_type(block->buffer))) {
+            host_pool_in_use = true;
+            break;
+        }
+    }
+    // Placement probing is only valid for unloaded tensors (the operator
+    // probe rejects loaded ones); ignored and optional-missing states never
+    // load, so they must not be counted as missing host capacity.
+    for (TensorState* state : states) {
+        if (state == nullptr || should_ignore(*state) || is_optional_missing_tensor(state->name)) {
+            continue;
+        }
+        if (!state->loaded_to_params_backend && state->tensor != nullptr &&
+            model_manager_buffer_type_is_host_pinned(params_buffer_type_for(*state))) {
+            host_missing_bytes += ggml_nbytes(state->tensor);
+            host_pool_in_use    = true;
+        }
+    }
+    if (missing > 0 && model_manager_buffer_type_is_host_pinned(
+                           ggml_backend_get_default_buffer_type(request.compute_backend))) {
+        host_missing_bytes += missing;
+        host_pool_in_use    = true;
+    }
+    if (host_pool_in_use) {
+        // The +1 byte enforces the safety floor itself — once the host-pinned
+        // pool is in use, a request that pins nothing new must still fail
+        // when availability has fallen to zero, so reclamation runs before
+        // the driver's allocations start failing.
+        result.required_host_bytes = host_missing_bytes + 1;
+#ifdef _WIN32
+        MEMORYSTATUSEX status{};
+        status.dwLength = sizeof(status);
+        if (GlobalMemoryStatusEx(&status)) {
+            // Two thresholds, from measurements on the #1976 machine (32 GB
+            // RAM, both resident models loaded): the driver fails allocations
+            // once available physical memory drops below ~0.5-0.8 GB (512 MiB
+            // + request bytes = hard refusal region), while up to ~0.5 GB is
+            // consumed between a capacity check and the following submit
+            // (1 GiB = proactive reclamation watermark so eviction keeps the
+            // pool above the failing region instead of only reacting at it).
+            constexpr uint64_t host_safety_margin      = 512ULL * 1024ULL * 1024ULL;
+            constexpr uint64_t host_reclaim_watermark  = 1ULL * 1024ULL * 1024ULL * 1024ULL;
+            result.available_host_bytes               = status.ullAvailPhys > host_safety_margin
+                                                              ? static_cast<size_t>(status.ullAvailPhys - host_safety_margin)
+                                                              : 0;
+            if (status.ullAvailPhys < host_reclaim_watermark) {
+                result.requires_reclamation = true;
+            }
+        }
+#endif
+    }
+    LOG_VERBOSE("capacity check verdict on %s: required device %.2f MB / available %.2f MB, required budget %.2f MB / available %.2f MB, required host %.2f MB / available %.2f MB, fits %s",
+                ggml_backend_name(request.compute_backend),
+                result.required_device_bytes / (1024.0 * 1024.0),
+                result.available_device_bytes / (1024.0 * 1024.0),
+                result.required_budget_bytes / (1024.0 * 1024.0),
+                result.available_budget_bytes / (1024.0 * 1024.0),
+                result.required_host_bytes / (1024.0 * 1024.0),
+                result.available_host_bytes / (1024.0 * 1024.0),
+                result.fits() ? "yes" : "no");
     return result;
 }
 
@@ -1974,17 +2074,37 @@ bool ModelManager::ensure_compute_backend_capacity(
             }
         }
     };
+    // Offloaded parameter storage can also sit in host-pinned memory (the
+    // compute device's host buffer type); under host pressure those blocks
+    // must be reclaimable too, or only device staging gets evicted while the
+    // host pool cannot recover.
+    std::unordered_set<TensorState*> host_pinned_states;
+    for (const auto& block : params_storage_blocks_) {
+        if (block != nullptr && block->buffer != nullptr &&
+            model_manager_buffer_type_is_host_pinned(ggml_backend_buffer_get_type(block->buffer))) {
+            host_pinned_states.insert(block->states.begin(), block->states.end());
+        }
+    }
     auto can_demote_to_disk = [&](TensorState* state) {
-        return state != nullptr && state->tensor != nullptr && state->tensor->view_src == nullptr &&
-               state->residency_mode == ResidencyMode::ParamBackend &&
-               state->params_backend != nullptr && state->params_backend == state->compute_backend &&
-               state->compute_backend == compute_backend && !sd_backend_is_cpu(state->params_backend) &&
-               state->pin_count == 0 && !state->staged_to_compute_backend && state->loaded_to_params_backend &&
-               protected_states.find(state) == protected_states.end() &&
-               required_set.find(state) == required_set.end() &&
-               mapped_states.find(state) == mapped_states.end() &&
-               !state->sources.empty() && state->applied_lora_epoch == UINT64_MAX &&
-               state->split_buffer_type == nullptr;
+        if (state == nullptr || state->tensor == nullptr || state->tensor->view_src != nullptr ||
+            state->residency_mode != ResidencyMode::ParamBackend ||
+            state->pin_count != 0 || state->staged_to_compute_backend || !state->loaded_to_params_backend ||
+            protected_states.find(state) != protected_states.end() ||
+            required_set.find(state) != required_set.end() ||
+            state->sources.empty() || state->applied_lora_epoch != UINT64_MAX ||
+            state->split_buffer_type != nullptr) {
+            return false;
+        }
+        const bool device_resident = state->params_backend != nullptr &&
+                                     state->params_backend == state->compute_backend &&
+                                     state->compute_backend == compute_backend &&
+                                     !sd_backend_is_cpu(state->params_backend) &&
+                                     mapped_states.find(state) == mapped_states.end();
+        // Host-pinned parameter storage of an offloaded model reloads from
+        // the model loader like Disk-resident states; mmap-backed blocks are
+        // still freed (the pinned copy is a separate allocation).
+        const bool host_pinned = host_pinned_states.find(state) != host_pinned_states.end();
+        return device_resident || host_pinned;
     };
     for (TensorState* state : global_candidates) {
         if (!can_demote_to_disk(state)) {
@@ -1993,14 +2113,23 @@ bool ModelManager::ensure_compute_backend_capacity(
         state->residency_mode = ResidencyMode::Disk;
         demoted_states.push_back(state);
         add_evictable_state(state);
+        if (host_pinned_states.find(state) != host_pinned_states.end()) {
+            // Host-pinned offloaded params are resident through their
+            // parameter block (loaded_to_params_backend), not through
+            // compute staging, so add_evictable_state's residency test
+            // cannot admit them — insert directly so the block becomes
+            // releasable.
+            eviction_states.insert(state);
+        }
         if (release_eviction_states()) {
             revert_demotions();
             return true;
         }
     }
     revert_demotions();
 
-    if (check_capacity(request, required_states).fits()) {
+    const auto final_capacity = check_capacity(request, required_states);
+    if (!final_capacity.requires_reclamation && final_capacity.fits()) {
         return true;
     }
     const auto capacity                = check_capacity(request, required_states, true);
@@ -2010,11 +2139,19 @@ bool ModelManager::ensure_compute_backend_capacity(
     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",
+    const std::string required_host    = capacity.required_host_bytes == 0
+                                             ? "none"
+                                             : sd_format("%.2f MB", capacity.required_host_bytes / (1024.0 * 1024.0));
+    const std::string available_host   = capacity.available_host_bytes == SIZE_MAX
+                                             ? "unknown"
+                                             : sd_format("%.2f MB", capacity.available_host_bytes / (1024.0 * 1024.0));
+    LOG_WARN("model manager cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget / %s host, available %s device / %s budget / %s host",
              ggml_backend_name(compute_backend),
              capacity.required_device_bytes / (1024.0 * 1024.0),
              capacity.required_budget_bytes / (1024.0 * 1024.0),
+             required_host.c_str(),
              available_device.c_str(),
-             available_budget.c_str());
+             available_budget.c_str(),
+             available_host.c_str());
     return false;
 }

src/model_manager.h

@@ -150,12 +150,15 @@ class ModelManager : public DeviceResidencyManager {
     struct CapacityCheck {
         size_t required_device_bytes  = 0;
         size_t required_budget_bytes  = 0;
+        size_t required_host_bytes    = 0;
         size_t available_device_bytes = SIZE_MAX;
         size_t available_budget_bytes = SIZE_MAX;
+        size_t available_host_bytes   = SIZE_MAX;
         bool requires_reclamation     = false;
         bool fits() const {
             return required_device_bytes <= available_device_bytes &&
-                   required_budget_bytes <= available_budget_bytes;
+                   required_budget_bytes <= available_budget_bytes &&
+                   required_host_bytes <= available_host_bytes;
         }
     };
     CapacityCheck check_capacity(const DeviceMemoryRequest& request,