fix: reserve 128 MiB headroom when selecting monolithic execution (#2046)

This commit is contained in:
leejet 2026-09-25 01:53:11 +08:00 committed by GitHub
parent 740c7ae193
commit 1a2330de68
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 16 additions and 1 deletions

View File

@ -63,6 +63,12 @@ See [backend selection](./backend.md) for full syntax.
When a graph has cut markers and its missing weights plus incremental compute workspace exceed the available device headroom, it runs its fixed segment list in order. A reusable monolithic compute buffer is not counted as a new allocation. An explicit `--max-vram` budget deducts already-resident managed weights and compute/cache buffers registered by every runner sharing the device, so later graph runs remain segmented when the full graph exceeds the budget. The current segment's weights are pinned during compute, and the next parameter-bearing segment is prefetched when the device supports asynchronous transfer. No opt-in streaming flag is required. When a graph has cut markers and its missing weights plus incremental compute workspace exceed the available device headroom, it runs its fixed segment list in order. A reusable monolithic compute buffer is not counted as a new allocation. An explicit `--max-vram` budget deducts already-resident managed weights and compute/cache buffers registered by every runner sharing the device, so later graph runs remain segmented when the full graph exceeds the budget. The current segment's weights are pinned during compute, and the next parameter-bearing segment is prefetched when the device supports asynchronous transfer. No opt-in streaming flag is required.
When choosing between monolithic and segmented execution, the runner requires
an additional 128 MiB of headroom in both available device memory and any explicit
managed budget. This planning headroom absorbs small allocation estimate changes;
subsequent capacity checks can consume it while still preserving the 512 MiB device
scratch reserve and respecting the managed budget.
- `--max-vram <GiB>` optionally lowers the live-memory limit. A positive value is a managed per-device budget, `0` uses the device's current free memory without an explicit budget, and a negative value snapshots free memory at startup while reserving that many GiB (`--max-vram -1` reserves about 1 GiB). Driver contexts and unrelated external allocations remain outside the managed budget. - `--max-vram <GiB>` optionally lowers the live-memory limit. A positive value is a managed per-device budget, `0` uses the device's current free memory without an explicit budget, and a negative value snapshots free memory at startup while reserving that many GiB (`--max-vram -1` reserves about 1 GiB). Driver contexts and unrelated external allocations remain outside the managed budget.
- `--disable-prefetch` disables asynchronous next-segment prefetch while retaining synchronous loading, eviction, and segmented execution. - `--disable-prefetch` disables asynchronous next-segment prefetch while retaining synchronous loading, eviction, and segmented execution.
- `--disable-segmented-compute` forces monolithic graph execution for diagnostics or compatibility, even when the automatic memory check would select segments. - `--disable-segmented-compute` forces monolithic graph execution for diagnostics or compatibility, even when the automatic memory check would select segments.

View File

@ -837,11 +837,20 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
last_compute_status_ = GGML_STATUS_ALLOC_FAILED; last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
return std::nullopt; return std::nullopt;
} }
auto fits_monolithic = [&]() {
// Planning headroom absorbs allocation estimate drift; execution keeps the normal limits.
constexpr size_t planning_headroom = 128ULL * 1024ULL * 1024ULL;
auto requests = memory_requests(full_measurement.buffers, cache_.pending_bytes(graph));
for (auto& request : requests) {
request.pending_allocation_bytes = add_bytes(request.pending_allocation_bytes, planning_headroom);
}
return fits(requests, params);
};
auto manager = residency_manager.lock(); auto manager = residency_manager.lock();
const bool segmented = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) && const bool segmented = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) &&
manager != nullptr && manager->segmented_compute_enabled() && manager != nullptr && manager->segmented_compute_enabled() &&
cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1 && cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1 &&
!fits(memory_requests(full_measurement.buffers, cache_.pending_bytes(graph)), params); !fits_monolithic();
ggml_graph_cut::Plan monolithic_plan; ggml_graph_cut::Plan monolithic_plan;
if (!segmented) { if (!segmented) {
monolithic_plan.segments.emplace_back(); monolithic_plan.segments.emplace_back();