From 462d675018dee398881ab50a86df58aaaa5aab40 Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 6 Sep 2026 22:30:45 +0800 Subject: [PATCH] refactor: unify runner lifecycles and weight residency (#1940) --- .gitignore | 1 + docs/animatediff.md | 2 +- docs/backend.md | 25 +- docs/performance.md | 39 +- docs/sd.md | 2 +- docs/sefi_image.md | 2 +- examples/common/common.cpp | 16 +- examples/common/common.h | 14 +- examples/common/log.cpp | 1 + include/stable-diffusion.h | 6 +- src/conditioning/conditioner.hpp | 192 +--- src/core/compute_workspace.cpp | 257 +++++ src/core/compute_workspace.h | 64 ++ src/core/ggml_extend.hpp | 1033 +++----------------- src/core/ggml_graph_cut.cpp | 574 +++++------ src/core/ggml_graph_cut.h | 54 +- src/core/ggml_runner.cpp | 296 ++++++ src/core/layer_split_partition.cpp | 16 +- src/core/layer_stream_prefetch.cpp | 124 --- src/core/layer_stream_prefetch.h | 48 - src/core/runner_cache.cpp | 211 ++++ src/core/runner_cache.h | 66 ++ src/core/segment_graph_bindings.cpp | 137 +++ src/core/segment_graph_bindings.h | 52 + src/core/segment_weight_pipeline.cpp | 205 ++++ src/core/segment_weight_pipeline.h | 62 ++ src/core/tensor_ggml.hpp | 25 +- src/detailer.cpp | 2 +- src/device_residency_manager.h | 76 ++ src/extensions/generation_extension.h | 2 +- src/extensions/photomaker_extension.cpp | 4 +- src/model/adapter/ip_adapter.hpp | 2 +- src/model/adapter/lora.hpp | 26 +- src/model/adapter/pmid.hpp | 2 +- src/model/detector/yolov8.h | 2 +- src/model/diffusion/anima.hpp | 2 +- src/model/diffusion/boogu.hpp | 2 +- src/model/diffusion/control.hpp | 29 +- src/model/diffusion/ernie_image.hpp | 2 +- src/model/diffusion/flux.hpp | 2 +- src/model/diffusion/hidream_o1.hpp | 14 +- src/model/diffusion/hunyuan.hpp | 2 +- src/model/diffusion/ideogram4.hpp | 2 +- src/model/diffusion/krea2.hpp | 2 +- src/model/diffusion/lens.hpp | 2 +- src/model/diffusion/lingbot_video.hpp | 2 +- src/model/diffusion/ltxv.hpp | 2 +- src/model/diffusion/mage_flow.hpp | 2 +- src/model/diffusion/minimax_h3.hpp | 2 - src/model/diffusion/minit2i.hpp | 2 +- src/model/diffusion/mmdit.hpp | 2 +- src/model/diffusion/pid.hpp | 2 +- src/model/diffusion/qwen_image.hpp | 2 +- src/model/diffusion/unet.hpp | 2 +- src/model/diffusion/wan.hpp | 2 +- src/model/diffusion/z_image.hpp | 2 +- src/model/te/clip.hpp | 6 +- src/model/te/llm.hpp | 28 +- src/model/te/t5.hpp | 6 +- src/model/upscaler/esrgan.hpp | 2 +- src/model/upscaler/ltx_latent_upscaler.hpp | 2 +- src/model/vae/auto_encoder_kl.hpp | 2 +- src/model/vae/hunyuan_vae.hpp | 4 +- src/model/vae/ltx_audio_vae.hpp | 2 +- src/model/vae/ltx_vae.hpp | 8 +- src/model/vae/mage_vae.hpp | 2 +- src/model/vae/minimax_h3_audio_vae.hpp | 4 +- src/model/vae/minimax_h3_vae.hpp | 2 - src/model/vae/tae.hpp | 4 +- src/model/vae/vae.hpp | 4 +- src/model/vae/wan_vae.hpp | 6 +- src/model_manager.cpp | 689 +++++++++++-- src/model_manager.h | 97 +- src/model_manager_prefetch.cpp | 258 +++-- src/stable-diffusion.cpp | 130 ++- src/upscaler.cpp | 10 +- src/upscaler.h | 2 - src/weight_manager.h | 25 - 78 files changed, 2963 insertions(+), 2051 deletions(-) create mode 100644 src/core/compute_workspace.cpp create mode 100644 src/core/compute_workspace.h create mode 100644 src/core/ggml_runner.cpp delete mode 100644 src/core/layer_stream_prefetch.cpp delete mode 100644 src/core/layer_stream_prefetch.h create mode 100644 src/core/runner_cache.cpp create mode 100644 src/core/runner_cache.h create mode 100644 src/core/segment_graph_bindings.cpp create mode 100644 src/core/segment_graph_bindings.h create mode 100644 src/core/segment_weight_pipeline.cpp create mode 100644 src/core/segment_weight_pipeline.h create mode 100644 src/device_residency_manager.h delete mode 100644 src/weight_manager.h diff --git a/.gitignore b/.gitignore index f828af59..80599b3f 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ build*/ cmake-build-*/ test/ +tests/ .vscode/ .idea/ .cache/ diff --git a/docs/animatediff.md b/docs/animatediff.md index cf71a074..2e923b0b 100644 --- a/docs/animatediff.md +++ b/docs/animatediff.md @@ -79,7 +79,7 @@ Low-VRAM streaming (verified with a 2 GiB cap on RTX 3060): .\bin\Release\sd-cli.exe -M vid_gen \ --model ..\models\checkpoints\realisticVisionV60B1.safetensors \ --motion-module ..\models\animatediff\mm_sd15_v3.safetensors \ - --max-vram 2.0 --stream-layers --diffusion-fa \ + --max-vram 2.0 --diffusion-fa \ -p "photo of coastline, rocks, storm weather, wind, waves, lightning" \ --cfg-scale 8.0 --sampling-method euler --scheduler discrete \ -H 384 -W 384 --video-frames 8 --fps 8 --steps 20 -s 42 \ diff --git a/docs/backend.md b/docs/backend.md index c936b223..29778b1d 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -41,7 +41,11 @@ sd-cli -m model.safetensors -p "a cat" --backend cuda0 --params-backend disk sd-cli -m model.safetensors -p "a cat" --backend diffusion=cuda0,vae=vulkan0 --max-vram cuda0=6,vulkan0=2 ``` -The budget applies to every module running on that backend. +The value is a shared per-device budget for managed weights and registered +runner compute/cache buffers. Live free memory can lower the effective limit +for each graph run. Driver contexts and allocations made outside the managed +model runners are not part of this accounting, so it is not a hard physical +VRAM cap. Module names are case-insensitive. Hyphens and underscores in module names are ignored, so `clip_vision`, `clip-vision`, and `clipvision` are equivalent. @@ -79,9 +83,10 @@ with `--params-backend diffusion=disk`, released directly from) its own device; an explicit assignment such as `te=cpu` keeps the parameters on that backend and stages each range to its device on demand. -Layer split cannot be combined with `--max-vram` graph-cut segmentation or -`--stream-layers` for the split module; those are single-device mechanisms and -are disabled for it. +Layer split uses the fixed graph-cut plan to assign blocks across devices, but +single-device segmented execution and next-segment prefetch are disabled for +the split module. `--max-vram` can still provide the per-device limits used by +layer split and auto-fit. Use `--list-devices` to see the device names available on the system. @@ -104,12 +109,18 @@ Compared to a layer split this uses all GPUs within every layer (instead of sequentially device by device) at the cost of a cross-device reduction per matmul - usually the faster option when the devices have fast interconnect. -Row split requires backend support for split buffers and is currently -available on CUDA only; on other backends (or when the listed devices belong -to different backend registries) the module falls back to a layer split. +Row split requires a compatible split-buffer export from the linked GGML +backend. If it is unavailable (or the listed devices belong to different backend +registries), the module falls back to a layer split. Embeddings, normalization weights, biases and other non-block tensors stay in regular buffers on the main device. +Row-split execution can use graph segments, but split weights are loaded +synchronously instead of using the normal single-device prefetch path. Because +GGML does not expose exact shard allocation sizes, the managed budget currently +counts a split buffer's full size on each participating device. This is a +conservative bound and can reject otherwise feasible layouts. + Direct ("immediately") LoRA application cannot patch row-split tensors; with `--split-mode row` the automatic LoRA mode selects runtime application, and an explicit `--lora-apply-mode immediately` skips the split tensors with a diff --git a/docs/performance.md b/docs/performance.md index f5d39fa8..0aa0afbe 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -14,8 +14,12 @@ Run by adding `--diffusion-fa` to the arguments and watch for: ``` and the compute buffer shrink in the debug log: ``` -[DEBUG] ggml_extend.hpp:1004 - flux compute buffer size: 650.00 MB(VRAM) +[DEBUG] ggml_runner.cpp:280 - flux compute buffer size: 650.00 MB(VRAM) on CUDA0 (peak across 1 segment) ``` +This reports the actual peak compute workspace capacity per backend, including +CPU fallback. It excludes weights and cache buffers. Within a runner lifecycle, +the summary is printed only on the first graph or when backend capacities or the +segment count change. ## Offload weights to the CPU to save VRAM without reducing generation speed. @@ -43,7 +47,7 @@ Use disk params to reduce both VRAM and RAM usage: --backend cuda0 --params-backend disk ``` -This reloads parameters from the model file on demand and releases them after use. It has the lowest memory residency, but can be slower because weights must be read again. `disk` is never selected implicitly; set it explicitly when RAM usage matters more than reload cost. +This reloads parameters from the model file on demand, retains unpinned compute copies while space permits, and releases them under pressure or at module-run completion. It has the lowest source-memory residency, but can be slower because evicted weights must be read again. `disk` is never selected implicitly; set it explicitly when RAM usage matters more than reload cost. Per-module assignments can target only the largest modules: @@ -53,26 +57,37 @@ Per-module assignments can target only the largest modules: See [backend selection](./backend.md) for full syntax. -## Run models that don't fit in VRAM (CPU streaming). +## Run models that don't fit in VRAM (automatic segmented execution). -`--offload-to-cpu` alone keeps every parameter in system RAM and stages it to the runtime backend on first use, then leaves it resident there. If the diffusion model is larger than the runtime backend's free memory (e.g. Flux dev at bf16 on an 8 GiB GPU), that residency stops fitting during the sampling loop and generation fails. Two additional flags make it fit by trading a small amount of speed for room: +`--offload-to-cpu` keeps the source parameters in system RAM and creates compute-side GPU replicas on demand. Unpinned replicas remain resident for reuse, but automatic graph-cut execution evicts them from the last segment backward when the next weight or compute allocation needs space. Disk-backed parameters follow the same policy without retaining a RAM source copy. -- `--max-vram ` sets a VRAM budget the graph-cut segmenter respects. It cuts each forward pass into segments sized to fit the budget, running them in sequence and freeing intermediate activations between them. Negative values auto-detect free VRAM and spare the given amount (`--max-vram -1` uses most of the free VRAM and keeps ~1 GiB headroom), a positive value caps the budget, `0` disables segmentation. -- `--stream-layers` streams the diffusion model's transformer blocks one at a time. While one block computes, the next block's parameters are prefetched automatically from the CPU on a separate transfer queue; parameters are evicted when the residency budget is reached. This flag only takes effect when the diffusion params backend is CPU, so it must be combined with `--offload-to-cpu` (or an explicit `--params-backend diffusion=cpu`); a warning is logged and the flag is ignored otherwise. -- `--disable-prefetch` disables the asynchronous next-block prefetch while retaining synchronous `--stream-layers` execution. This is mainly useful for debugging or backends where transfer and compute do not overlap effectively. +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. -The three flags stack. The recommended shape for "biggest model my card can host": +- `--max-vram ` 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-segmented-compute` forces monolithic graph execution for diagnostics or compatibility, even when the automatic memory check would select segments. + +Single-device monolithic execution also reclaims unpinned weight replicas before +loading weights or allocating compute workspace, including graphs without cut +markers and runs with `--disable-segmented-compute`. It still respects the managed +device budget and fails if the graph cannot fit after reclamation. + +Segment completion releases active workspace use while retaining the runner's +allocator/scheduler capacity. Compatible gallocr reservations are reused across +graphs; idle workspaces can be reclaimed under pressure and are freed at runner +completion. Cross-graph caches survive individual graphs, but cut buffers do not. + +The recommended shape for "biggest model my card can host" is: ```shell sd-cli --diffusion-model flux1-dev.safetensors ... \ - --offload-to-cpu --max-vram -1 --stream-layers + --offload-to-cpu --max-vram -1 ``` - `--offload-to-cpu`: params in RAM, staged as needed. -- `--max-vram -1`: use most of the free VRAM as the compute budget, spare 1 GiB headroom, let the graph-cut segmenter split each forward pass to fit. -- `--stream-layers`: on top of the segmenter, stream individual transformer blocks so their weights don't all need to be resident at once. +- `--max-vram -1`: reserve about 1 GiB from the startup free-memory snapshot; live free memory can still lower the effective limit for every graph. -Ordered from fastest to smallest-VRAM: no flags → `--offload-to-cpu` → `--offload-to-cpu --max-vram ` → `--offload-to-cpu --max-vram --stream-layers`. Each step down costs a few percent of throughput to buy more room; combined they can run models roughly 3-4x larger than the raw VRAM would allow. +Use `--params-backend diffusion=disk` instead when reducing system RAM residency is more important than avoiding repeated model-file reads. ## Use quantization to reduce memory usage. diff --git a/docs/sd.md b/docs/sd.md index 2ba331ef..44a4c1f9 100644 --- a/docs/sd.md +++ b/docs/sd.md @@ -3,7 +3,7 @@ - download original weights(.ckpt or .safetensors). For example - Stable Diffusion v1.4 from https://huggingface.co/CompVis/stable-diffusion-v-1-4-original - Stable Diffusion v1.5 from https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5 - - Stable Diffuison v2.1 from https://huggingface.co/Manojb/stable-diffusion-2-1-base + - Stable Diffusion v2.1 from https://huggingface.co/Manojb/stable-diffusion-2-1-base - Stable Diffusion 3 2B from https://huggingface.co/stabilityai/stable-diffusion-3-medium ### txt2img example diff --git a/docs/sefi_image.md b/docs/sefi_image.md index 98e126ac..c6a3587b 100644 --- a/docs/sefi_image.md +++ b/docs/sefi_image.md @@ -44,7 +44,7 @@ The dispatcher picks `alpha` from the filename (`turbo` substring => 1.0, otherw ### 5B (needs streaming on 12 GiB VRAM) ``` -./build/bin/sd-cli --diffusion-model /path/to/sefi_5b_turbo.safetensors --vae /path/to/flux2_ae.safetensors --llm /path/to/qwen3_vl_4b.safetensors -p "a photograph of an orange tabby cat sitting on a couch" --cfg-scale 1.0 --steps 4 -W 1024 -H 1024 -s 42 --diffusion-fa --max-vram 8 --stream-layers --offload-to-cpu -o out.png +./build/bin/sd-cli --diffusion-model /path/to/sefi_5b_turbo.safetensors --vae /path/to/flux2_ae.safetensors --llm /path/to/qwen3_vl_4b.safetensors -p "a photograph of an orange tabby cat sitting on a couch" --cfg-scale 1.0 --steps 4 -W 1024 -H 1024 -s 42 --diffusion-fa --max-vram 8 --offload-to-cpu -o out.png ``` SeFi-Image 5B turbo example diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 2053b654..aa18f499 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -502,7 +502,7 @@ ArgOptions SDContextParams::get_options() { &rpc_servers}, {"", "--max-vram", - "maximum VRAM budget in GiB for graph-cut segmented execution. Accepts a single value or assignments by backend/device, e.g. 6 or cuda0=6,vulkan0=4. 0 disables graph splitting; a negative value auto-detects free VRAM, sparing the specified value", + "optional per-device budget in GiB for managed weights and runner buffers during automatic graph-cut execution. Accepts a single value or assignments by backend/device, e.g. 6 or cuda0=6,vulkan0=4. 0 uses live free VRAM without an explicit budget; a negative value reserves that much free VRAM", 0, &max_vram}, }; @@ -516,14 +516,14 @@ ArgOptions SDContextParams::get_options() { }; options.bool_options = { - {"", - "--stream-layers", - "enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram; defaults to false)", - true, &stream_layers}, {"", "--disable-prefetch", - "disable asynchronous layer prefetch while keeping synchronous --stream-layers behavior (defaults to false)", + "disable asynchronous next-segment weight prefetch (defaults to false)", true, &disable_prefetch}, + {"", + "--disable-segmented-compute", + "force monolithic graph execution even when automatic graph cutting is needed (defaults to false)", + true, &disable_segmented_compute}, {"", "--eager-load", "load all params into the params backend at model-load time instead of lazily on first use (defaults to false)", @@ -835,8 +835,8 @@ std::string SDContextParams::to_string() const { << " sampler_rng_type: " << sd_rng_type_name(sampler_rng_type) << ",\n" << " offload_params_to_cpu: " << (offload_params_to_cpu ? "true" : "false") << ",\n" << " max_vram: \"" << max_vram << "\",\n" - << " stream_layers: " << (stream_layers ? "true" : "false") << ",\n" << " disable_prefetch: " << (disable_prefetch ? "true" : "false") << ",\n" + << " disable_segmented_compute: " << (disable_segmented_compute ? "true" : "false") << ",\n" << " eager_load: " << (eager_load ? "true" : "false") << ",\n" << " backend: \"" << backend << "\",\n" << " params_backend: \"" << params_backend << "\",\n" @@ -908,8 +908,8 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { sd_ctx_params.force_sdxl_vae_conv_scale = force_sdxl_vae_conv_scale; sd_ctx_params.vae_format = str_to_vae_format(vae_format); sd_ctx_params.max_vram = max_vram.c_str(); - sd_ctx_params.stream_layers = stream_layers; sd_ctx_params.disable_prefetch = disable_prefetch; + sd_ctx_params.disable_segmented_compute = disable_segmented_compute; sd_ctx_params.eager_load = eager_load; sd_ctx_params.backend = effective_backend.c_str(); sd_ctx_params.params_backend = effective_params_backend.c_str(); diff --git a/examples/common/common.h b/examples/common/common.h index 0e726fab..83185d37 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -146,13 +146,13 @@ struct SDContextParams { std::map embedding_map; std::vector embedding_vec; - rng_type_t rng_type = CUDA_RNG; - rng_type_t sampler_rng_type = RNG_TYPE_COUNT; - bool offload_params_to_cpu = false; - std::string max_vram = "0"; - bool stream_layers = false; - bool disable_prefetch = false; - bool eager_load = false; + rng_type_t rng_type = CUDA_RNG; + rng_type_t sampler_rng_type = RNG_TYPE_COUNT; + bool offload_params_to_cpu = false; + std::string max_vram = "0"; + bool disable_prefetch = false; + bool disable_segmented_compute = false; + bool eager_load = false; std::string backend; std::string params_backend; std::string split_mode; diff --git a/examples/common/log.cpp b/examples/common/log.cpp index 2c434391..61988ad8 100644 --- a/examples/common/log.cpp +++ b/examples/common/log.cpp @@ -88,6 +88,7 @@ void log_print(enum sd_log_level_t level, const char* log, bool verbose, bool co } else { fprintf(out_stream, "[%-5s] ", level_str); } + fflush(out_stream); print_utf8(out_stream, log); fflush(out_stream); } diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 9c85114f..ff4e3063 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -229,9 +229,8 @@ typedef struct { bool vae_conv_direct; bool force_sdxl_vae_conv_scale; enum sd_vae_format_t vae_format; - const char* max_vram; // GiB budget or backend assignment spec for graph-cut segmented param offload (0 = disabled, -1 = auto) - bool stream_layers; // Enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram) - bool disable_prefetch; // Disable asynchronous layer prefetch while retaining synchronous stream_layers behavior + const char* max_vram; // Optional per-device GiB budget for managed weights and runner buffers; 0 uses live free VRAM without an explicit budget + bool disable_prefetch; // Disable asynchronous next-segment weight prefetch bool eager_load; // Load all params into the params backend at model-load time instead of lazily on first use const char* backend; const char* params_backend; @@ -239,6 +238,7 @@ typedef struct { bool auto_fit; const char* rpc_servers; const char* model_args; + bool disable_segmented_compute; // Force monolithic graph execution even when automatic graph cutting would fit memory better } sd_ctx_params_t; typedef struct { diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index 8968676b..db28baab 100644 --- a/src/conditioning/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -142,14 +142,13 @@ public: virtual void get_param_tensors(std::map& tensors) = 0; virtual void get_param_tensor_ops(std::map& tensor_ops) {} virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {} - virtual void set_stream_layers_enabled(bool enabled) {} virtual void set_runtime_backends(const std::vector& backends) {} virtual void set_graph_cut_layer_split_enabled(bool enabled) {} virtual void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) {} virtual void get_layer_split_param_tensors(std::map& tensors) {} virtual void set_flash_attention_enabled(bool enabled) = 0; virtual void set_weight_adapter(const std::shared_ptr& adapter) {} - virtual void runner_done() {} + virtual void runner_end() {} }; // ldm.modules.encoders.modules.FrozenCLIPEmbedder @@ -202,13 +201,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { } } - void set_stream_layers_enabled(bool enabled) override { - text_model->set_stream_layers_enabled(enabled); - if (sd_version_is_sdxl(version)) { - text_model2->set_stream_layers_enabled(enabled); - } - } - void set_runtime_backends(const std::vector& backends) override { text_model->set_runtime_backends(backends); if (sd_version_is_sdxl(version)) { @@ -244,10 +236,10 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { } } - void runner_done() override { - text_model->runner_done(); + void runner_end() override { + text_model->runner_end(); if (sd_version_is_sdxl(version)) { - text_model2->runner_done(); + text_model2->runner_end(); } } @@ -461,9 +453,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { max_token_idx, false, clip_skip, - false, - true, - true); + false); GGML_ASSERT(!chunk_hidden_states.empty()); if (sd_version_is_sdxl(version)) { auto chunk_hidden_states2 = text_model2->compute(n_threads, @@ -473,9 +463,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { max_token_idx, false, clip_skip, - false, - true, - true); + false); GGML_ASSERT(!chunk_hidden_states2.empty()); chunk_hidden_states = sd::ops::concat(chunk_hidden_states, chunk_hidden_states2, 0); @@ -487,9 +475,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { max_token_idx, true, clip_skip, - false, - true, - true); + false); GGML_ASSERT(!pooled.empty()); } } @@ -608,7 +594,7 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(pixel_values, return_pooled, clip_skip); }; - return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true, true, true)); + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true)); } }; @@ -675,18 +661,6 @@ struct SD3CLIPEmbedder : public Conditioner { } } - void set_stream_layers_enabled(bool enabled) override { - if (clip_l) { - clip_l->set_stream_layers_enabled(enabled); - } - if (clip_g) { - clip_g->set_stream_layers_enabled(enabled); - } - if (t5) { - t5->set_stream_layers_enabled(enabled); - } - } - void set_runtime_backends(const std::vector& backends) override { if (clip_l) { clip_l->set_runtime_backends(backends); @@ -753,15 +727,15 @@ struct SD3CLIPEmbedder : public Conditioner { } } - void runner_done() override { + void runner_end() override { if (clip_l) { - clip_l->runner_done(); + clip_l->runner_end(); } if (clip_g) { - clip_g->runner_done(); + clip_g->runner_end(); } if (t5) { - t5->runner_done(); + t5->runner_end(); } } @@ -881,9 +855,7 @@ struct SD3CLIPEmbedder : public Conditioner { max_token_idx, false, clip_skip, - false, - true, - true); + false); GGML_ASSERT(!chunk_hidden_states_l.empty()); chunk_hidden_states_l = ::apply_token_weights(std::move(chunk_hidden_states_l), chunk_weights); @@ -897,9 +869,7 @@ struct SD3CLIPEmbedder : public Conditioner { max_token_idx, true, clip_skip, - false, - true, - true); + false); GGML_ASSERT(!pooled_l.empty()); } } else { @@ -928,9 +898,7 @@ struct SD3CLIPEmbedder : public Conditioner { max_token_idx, false, clip_skip, - false, - true, - true); + false); GGML_ASSERT(!chunk_hidden_states_g.empty()); chunk_hidden_states_g = ::apply_token_weights(std::move(chunk_hidden_states_g), chunk_weights); @@ -944,9 +912,7 @@ struct SD3CLIPEmbedder : public Conditioner { max_token_idx, true, clip_skip, - false, - true, - true); + false); GGML_ASSERT(!pooled_g.empty()); } } else { @@ -969,9 +935,7 @@ struct SD3CLIPEmbedder : public Conditioner { chunk_hidden_states_t5 = t5->compute(n_threads, input_ids, sd::Tensor(), - false, - true, - true); + false); GGML_ASSERT(!chunk_hidden_states_t5.empty()); chunk_hidden_states_t5 = ::apply_token_weights(std::move(chunk_hidden_states_t5), chunk_weights); } else { @@ -1079,15 +1043,6 @@ struct FluxCLIPEmbedder : public Conditioner { } } - void set_stream_layers_enabled(bool enabled) override { - if (clip_l) { - clip_l->set_stream_layers_enabled(enabled); - } - if (t5) { - t5->set_stream_layers_enabled(enabled); - } - } - void set_runtime_backends(const std::vector& backends) override { if (clip_l) { clip_l->set_runtime_backends(backends); @@ -1139,12 +1094,12 @@ struct FluxCLIPEmbedder : public Conditioner { } } - void runner_done() override { + void runner_end() override { if (clip_l) { - clip_l->runner_done(); + clip_l->runner_end(); } if (t5) { - t5->runner_done(); + t5->runner_end(); } } @@ -1247,9 +1202,7 @@ struct FluxCLIPEmbedder : public Conditioner { max_token_idx, true, clip_skip, - false, - true, - true); + false); GGML_ASSERT(!pooled.empty()); } else { pooled = sd::Tensor::zeros({768}); @@ -1268,9 +1221,7 @@ struct FluxCLIPEmbedder : public Conditioner { chunk_hidden_states = t5->compute(n_threads, input_ids, sd::Tensor(), - false, - true, - true); + false); GGML_ASSERT(!chunk_hidden_states.empty()); chunk_hidden_states = ::apply_token_weights(std::move(chunk_hidden_states), chunk_weights); if (zero_out_masked) { @@ -1366,12 +1317,6 @@ struct T5CLIPEmbedder : public Conditioner { } } - void set_stream_layers_enabled(bool enabled) override { - if (t5) { - t5->set_stream_layers_enabled(enabled); - } - } - void set_runtime_backends(const std::vector& backends) override { if (t5) { t5->set_runtime_backends(backends); @@ -1408,9 +1353,9 @@ struct T5CLIPEmbedder : public Conditioner { } } - void runner_done() override { + void runner_end() override { if (t5) { - t5->runner_done(); + t5->runner_end(); } } @@ -1508,9 +1453,7 @@ struct T5CLIPEmbedder : public Conditioner { auto chunk_hidden_states = t5->compute(n_threads, input_ids, t5_attn_mask_chunk, - false, - true, - true); + false); GGML_ASSERT(!chunk_hidden_states.empty()); chunk_hidden_states = apply_token_weights(std::move(chunk_hidden_states), chunk_weights); @@ -1582,12 +1525,6 @@ struct MiniT2IConditioner : public Conditioner { } } - void set_stream_layers_enabled(bool enabled) override { - if (t5) { - t5->set_stream_layers_enabled(enabled); - } - } - void set_runtime_backends(const std::vector& backends) override { if (t5) { t5->set_runtime_backends(backends); @@ -1624,9 +1561,9 @@ struct MiniT2IConditioner : public Conditioner { } } - void runner_done() override { + void runner_end() override { if (t5) { - t5->runner_done(); + t5->runner_end(); } } @@ -1657,9 +1594,7 @@ struct MiniT2IConditioner : public Conditioner { sd::Tensor hidden_states = t5->compute(n_threads, input_ids, sd::Tensor::from_vector(t5_mask), - false, - true, - true); + false); GGML_ASSERT(!hidden_states.empty()); result.c_crossattn = std::move(hidden_states); result.c_vector = sd::Tensor::from_vector(mask); @@ -1696,10 +1631,6 @@ struct AnimaConditioner : public Conditioner { llm->set_max_graph_vram_bytes(max_vram_bytes); } - void set_stream_layers_enabled(bool enabled) override { - llm->set_stream_layers_enabled(enabled); - } - void set_runtime_backends(const std::vector& backends) override { llm->set_runtime_backends(backends); } @@ -1724,8 +1655,8 @@ struct AnimaConditioner : public Conditioner { llm->set_weight_adapter(adapter); } - void runner_done() override { - llm->runner_done(); + void runner_end() override { + llm->runner_end(); } std::tuple, std::vector, std::vector, std::vector> tokenize(std::string text) { @@ -1787,9 +1718,7 @@ struct AnimaConditioner : public Conditioner { {}, {}, false, - false, - true, - true); + false); GGML_ASSERT(!hidden_states.empty()); hidden_states = apply_token_weights(std::move(hidden_states), qwen_weights); auto t5_ids_tensor = sd::Tensor::from_vector(t5_tokens); @@ -1887,13 +1816,6 @@ struct LLMEmbedder : public Conditioner { } } - void set_stream_layers_enabled(bool enabled) override { - llm->set_stream_layers_enabled(enabled); - if (byt5) { - byt5->set_stream_layers_enabled(enabled); - } - } - void set_runtime_backends(const std::vector& backends) override { llm->set_runtime_backends(backends); if (byt5) { @@ -1942,12 +1864,12 @@ struct LLMEmbedder : public Conditioner { } } - void runner_done() override { + void runner_end() override { if (llm) { - llm->runner_done(); + llm->runner_end(); } if (byt5) { - byt5->runner_done(); + byt5->runner_end(); } } @@ -2051,8 +1973,6 @@ struct LLMEmbedder : public Conditioner { out_layers, false, false, - true, - true, deepstack_image_embeds, image_grids); GGML_ASSERT(!hidden_states.empty()); @@ -2220,9 +2140,7 @@ struct LLMEmbedder : public Conditioner { prompt += ": "; add_vision_outputs(llm->encode_image_outputs(n_threads, resized, - false, - true, - true), + false), static_cast(resized.shape()[1]) / patch_size, static_cast(resized.shape()[0]) / patch_size); continue; @@ -2250,9 +2168,7 @@ struct LLMEmbedder : public Conditioner { auto pair = sd::ops::concat(first.unsqueeze(2), second.unsqueeze(2), 2); add_vision_outputs(llm->encode_video_block_outputs(n_threads, pair, - false, - true, - true), + false), static_cast(first.shape()[1]) / patch_size, static_cast(first.shape()[0]) / patch_size); } @@ -2263,9 +2179,7 @@ struct LLMEmbedder : public Conditioner { prompt += ": "; add_vision_outputs(llm->encode_image_outputs(n_threads, resized, - false, - true, - true), + false), static_cast(resized.shape()[1]) / patch_size, static_cast(resized.shape()[0]) / patch_size); } @@ -2352,7 +2266,7 @@ struct LLMEmbedder : public Conditioner { LOG_DEBUG("resize LingBotVideo ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); auto resized_image = clip_preprocess(image, w_bar, h_bar); - auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true); + auto image_embed = llm->encode_image(n_threads, resized_image, false); GGML_ASSERT(!image_embed.empty()); std::string image_prefix = prompt + img_prompt + "<|vision_start|>"; @@ -2411,7 +2325,7 @@ struct LLMEmbedder : public Conditioner { auto resized_image = clip_preprocess(image, w_bar, h_bar); - auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true); + auto image_embed = llm->encode_image(n_threads, resized_image, false); GGML_ASSERT(!image_embed.empty()); image_embeds.emplace_back(image_embed_idx, image_embed); image_embed_idx += 1 + static_cast(image_embed.shape()[1]) + 6; @@ -2494,7 +2408,7 @@ struct LLMEmbedder : public Conditioner { LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); auto resized_image = clip_preprocess(image, w_bar, h_bar); - auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true); + auto image_embed = llm->encode_image(n_threads, resized_image, false); GGML_ASSERT(!image_embed.empty()); std::string image_prefix = prompt_prefix + img_prompt + "<|vision_start|>"; @@ -2562,7 +2476,7 @@ struct LLMEmbedder : public Conditioner { LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); auto resized_image = clip_preprocess(image, w_bar, h_bar); - auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true); + auto image_embed = llm->encode_image(n_threads, resized_image, false); GGML_ASSERT(!image_embed.empty()); std::string image_prefix = prompt + img_prompt + "Picture " + std::to_string(i + 1) + ": <|vision_start|>"; @@ -2625,7 +2539,7 @@ struct LLMEmbedder : public Conditioner { LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); auto resized_image = clip_preprocess(image, w_bar, h_bar); - auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true); + auto image_embed = llm->encode_image(n_threads, resized_image, false); GGML_ASSERT(!image_embed.empty()); image_embeds.emplace_back(image_embed_idx, image_embed); image_embed_idx += 1 + static_cast(image_embed.shape()[1]) + 6; @@ -2857,9 +2771,7 @@ struct LLMEmbedder : public Conditioner { auto byt5_hidden_states = byt5->compute(n_threads, input_ids, sd::Tensor(), - false, - true, - true); + false); GGML_ASSERT(!byt5_hidden_states.empty()); extra_hidden_states_vec.push_back(std::move(byt5_hidden_states)); } @@ -2960,13 +2872,11 @@ struct LTXAVTextProjectionRunner : public GGMLRunner { sd::Tensor compute(int n_threads, const sd::Tensor& x, - bool auto_free = true, - bool free_compute_buffer = true, - bool free_compute_params = true) { + bool auto_runner_end = true) { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); }; - return take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params)); + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end)); } }; @@ -3060,9 +2970,9 @@ struct LTXAVEmbedder : public Conditioner { projector->set_weight_adapter(adapter); } - void runner_done() override { - llm->runner_done(); - projector->runner_done(); + void runner_end() override { + llm->runner_end(); + projector->runner_end(); } std::tuple, std::vector, std::vector> tokenize(std::string text, @@ -3129,9 +3039,7 @@ struct LTXAVEmbedder : public Conditioner { {}, {}, true, - false, - true, - true); + false); GGML_ASSERT(!hidden_states.empty()); hidden_states = apply_token_weights(std::move(hidden_states), weights); @@ -3190,7 +3098,7 @@ struct LTXAVEmbedder : public Conditioner { } hidden_states.reshape_({kNumStates * kHiddenSize, valid_tokens}); - return projector->compute(n_threads, hidden_states, false, true, true); + return projector->compute(n_threads, hidden_states, false); } SDCondition get_learned_condition(int n_threads, diff --git a/src/core/compute_workspace.cpp b/src/core/compute_workspace.cpp new file mode 100644 index 00000000..a40e2dbd --- /dev/null +++ b/src/core/compute_workspace.cpp @@ -0,0 +1,257 @@ +#include "core/compute_workspace.h" + +#include +#include +#include +#include +#include + +#include "core/ggml_extend_backend.h" +#include "core/ggml_graph_cut.h" +#include "ggml-cpu.h" +#include "ggml/src/ggml-impl.h" + +namespace sd { + ComputeWorkspace::~ComputeWorkspace() { + segment_end(); + release(); + ggml_backend_free(cpu_backend_); + } + + void ComputeWorkspace::set_extra_backends(const std::vector& backends) { + if (extra_backends_ != backends) { + GGML_ASSERT(!active_); + release(); + extra_backends_ = backends; + } + } + + bool ComputeWorkspace::needs_scheduler(ggml_cgraph* graph) const { + if (!extra_backends_.empty()) { + return true; + } + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + if (!ggml_backend_supports_op(backend_, ggml_graph_node(graph, i))) { + return true; + } + } + return false; + } + + ggml_backend_sched_t ComputeWorkspace::make_scheduler(size_t graph_size) { + std::vector backends{backend_}; + backends.insert(backends.end(), extra_backends_.begin(), extra_backends_.end()); + if (!sd_backend_is_cpu(backend_)) { + if (cpu_backend_ == nullptr) { + cpu_backend_ = sd_backend_cpu_init(); + } + if (cpu_backend_ == nullptr) { + return nullptr; + } + backends.push_back(cpu_backend_); + } + std::vector bufts; + for (auto backend : backends) { + auto buft = backend == cpu_backend_ + ? ggml_backend_dev_host_buffer_type(ggml_backend_get_device(backend_)) + : nullptr; + bufts.push_back(buft != nullptr ? buft : ggml_backend_get_default_buffer_type(backend)); + } + return ggml_backend_sched_new(backends.data(), bufts.data(), static_cast(backends.size()), + graph_size, false, false); + } + + bool ComputeWorkspace::measurement_matches(ggml_cgraph* graph, const Measurement& measurement) const { + return measurement.scheduler == needs_scheduler(graph); + } + + bool ComputeWorkspace::prepare(const Measurement& measurement) { + GGML_ASSERT(!active_); + if (measurement.buffers.empty()) { + return false; + } + const bool grows = std::any_of(measurement.buffers.begin(), measurement.buffers.end(), + [&](const BackendBufferSize& size) { return size.bytes > bytes(size.backend); }); + if (measurement.scheduler != (scheduler_ != nullptr) || grows) { + release(); + } + return true; + } + + bool ComputeWorkspace::release_excess(const Measurement& measurement) { + return std::any_of(measurement.buffers.begin(), measurement.buffers.end(), + [&](const BackendBufferSize& size) { return bytes(size.backend) > size.bytes; }) && + release(); + } + + bool ComputeWorkspace::allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes) { + GGML_ASSERT(!active_); + const bool use_scheduler = needs_scheduler(graph); + if (use_scheduler) { + if (allocator_ != nullptr) { + release(); + } + const size_t capacity = static_cast(graph->n_nodes + graph->n_leafs) + 8; + if (scheduler_ == nullptr || capacity > scheduler_capacity_) { + release(); + scheduler_ = make_scheduler(capacity); + scheduler_capacity_ = capacity; + } + if (scheduler_ == nullptr) { + return false; + } + ggml_backend_sched_reset(scheduler_); + assign_nodes(scheduler_, graph); + // Scheduler allocation rewrites sources. Split the execution graph only once. + if (!ggml_backend_sched_alloc_graph(scheduler_, graph)) { + release(); + return false; + } + } else { + if (scheduler_ != nullptr) { + release(); + } + if (allocator_ == nullptr) { + allocator_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); + } + auto signature = ggml_graph_cut::graph_layout(graph, true); + if (signature != reservation_) { + if (!ggml_gallocr_reserve(allocator_, graph)) { + release(); + return false; + } + reservation_ = std::move(signature); + ++reservations_; + } + if (!ggml_gallocr_alloc_graph(allocator_, graph)) { + release(); + return false; + } + } + active_ = true; + return true; + } + + ComputeWorkspace::Measurement ComputeWorkspace::measure( + ggml_cgraph* graph, + size_t direct_bytes, + const std::function& external_backend, + const AssignNodes& assign_nodes) { + if (!needs_scheduler(graph)) { + return {{{backend_, direct_bytes}}, false}; + } + std::vector tensors; + std::unordered_set seen; + auto visit = [&](const ggml_tensor* tensor) { + if (tensor != nullptr && seen.insert(tensor).second) { + tensors.push_back(tensor); + } + }; + for (int i = 0; i < graph->n_nodes; ++i) { + visit(graph->nodes[i]); + } + for (int i = 0; i < graph->n_leafs; ++i) { + visit(graph->leafs[i]); + } + for (size_t i = 0; i < tensors.size(); ++i) { + visit(tensors[i]->view_src); + for (auto source : tensors[i]->src) { + visit(source); + } + } + const size_t graph_size = tensors.size() + 8; + auto context = ggml_init({tensors.size() * ggml_tensor_overhead() + ggml_graph_overhead_custom(graph_size, false), nullptr, true}); + if (context == nullptr) { + return {}; + } + std::unordered_map copies; + std::map external_buffers; + for (auto tensor : tensors) { + auto copy = ggml_dup_tensor(context, tensor); + *copy = *tensor; + copies[tensor] = copy; + } + for (const auto& entry : copies) { + auto source = entry.first; + auto copy = entry.second; + copy->view_src = source->view_src == nullptr ? nullptr : copies.at(source->view_src); + for (int i = 0; i < GGML_MAX_SRC; ++i) { + copy->src[i] = source->src[i] == nullptr ? nullptr : copies.at(source->src[i]); + } + auto external = external_backend(source); + if (external != nullptr && source->view_src == nullptr) { + auto& buffer = external_buffers[external]; + if (buffer == nullptr) { + buffer = ggml_backend_alloc_buffer(external, 0); + GGML_ASSERT(buffer != nullptr); + ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + } + copy->buffer = buffer; + copy->data = reinterpret_cast(static_cast(1)); + copy->extra = nullptr; + } + } + auto copy_graph = ggml_new_graph_custom(context, graph_size, false); + copy_graph->n_nodes = graph->n_nodes; + copy_graph->n_leafs = graph->n_leafs; + for (int i = 0; i < graph->n_nodes; ++i) { + copy_graph->nodes[i] = copies.at(graph->nodes[i]); + } + for (int i = 0; i < graph->n_leafs; ++i) { + copy_graph->leafs[i] = copies.at(graph->leafs[i]); + } + Measurement result; + result.scheduler = true; + auto scheduler = make_scheduler(graph_size); + if (scheduler != nullptr) { + assign_nodes(scheduler, copy_graph); + std::vector sizes(extra_backends_.size() + 2); + ggml_backend_sched_reserve_size(scheduler, copy_graph, sizes.data()); + result.buffers.push_back({backend_, sizes[0]}); + for (size_t i = 0; i < extra_backends_.size(); ++i) { + result.buffers.push_back({extra_backends_[i], sizes[i + 1]}); + } + ggml_backend_sched_free(scheduler); + } + for (const auto& entry : external_buffers) { + ggml_backend_buffer_free(entry.second); + } + ggml_free(context); + return result; + } + + void ComputeWorkspace::synchronize() const { + if (scheduler_ != nullptr) { + ggml_backend_sched_synchronize(scheduler_); + } else { + ggml_backend_synchronize(backend_); + } + } + + void ComputeWorkspace::segment_end() { + if (active_) { + synchronize(); + active_ = false; + } + } + + bool ComputeWorkspace::release() { + if (active_) { + return false; + } + ggml_gallocr_free(allocator_); + allocator_ = nullptr; + ggml_backend_sched_free(scheduler_); + scheduler_ = nullptr; + scheduler_capacity_ = 0; + reservation_.clear(); + return true; + } + + size_t ComputeWorkspace::bytes(ggml_backend_t backend) const { + if (scheduler_ != nullptr) { + return ggml_backend_sched_get_buffer_size(scheduler_, backend); + } + return allocator_ != nullptr && backend == backend_ ? ggml_gallocr_get_buffer_size(allocator_, 0) : 0; + } +} diff --git a/src/core/compute_workspace.h b/src/core/compute_workspace.h new file mode 100644 index 00000000..7e73c23e --- /dev/null +++ b/src/core/compute_workspace.h @@ -0,0 +1,64 @@ +#ifndef __SD_CORE_COMPUTE_WORKSPACE_H__ +#define __SD_CORE_COMPUTE_WORKSPACE_H__ + +#include +#include + +#include "ggml-alloc.h" +#include "ggml-backend.h" + +namespace sd { + struct BackendBufferSize { + ggml_backend_t backend = nullptr; + size_t bytes = 0; + }; + + class ComputeWorkspace { + ggml_backend_t backend_; + std::vector extra_backends_; + ggml_backend_t cpu_backend_ = nullptr; + ggml_gallocr_t allocator_ = nullptr; + ggml_backend_sched_t scheduler_ = nullptr; + size_t scheduler_capacity_ = 0; + std::vector reservation_; + bool active_ = false; + size_t reservations_ = 0; + + ggml_backend_sched_t make_scheduler(size_t graph_size); + bool needs_scheduler(ggml_cgraph* graph) const; + + public: + struct Measurement { + std::vector buffers; + bool scheduler = false; + }; + using AssignNodes = std::function; + + explicit ComputeWorkspace(ggml_backend_t backend) + : backend_(backend) {} + ~ComputeWorkspace(); + ComputeWorkspace(const ComputeWorkspace&) = delete; + ComputeWorkspace& operator=(const ComputeWorkspace&) = delete; + + void set_extra_backends(const std::vector& backends); + bool measurement_matches(ggml_cgraph* graph, const Measurement& measurement) const; + bool prepare(const Measurement& measurement); + bool release_excess(const Measurement& measurement); + bool allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes); + Measurement measure( + ggml_cgraph* graph, + size_t direct_bytes, + const std::function& external_backend, + const AssignNodes& assign_nodes); + void synchronize() const; + void segment_end(); + bool release(); + bool active() const { return active_; } + ggml_backend_sched_t scheduler() const { return scheduler_; } + ggml_backend_t cpu_backend() const { return cpu_backend_; } + size_t bytes(ggml_backend_t backend) const; + size_t reservation_count() const { return reservations_; } + }; +} + +#endif // __SD_CORE_COMPUTE_WORKSPACE_H__ diff --git a/src/core/ggml_extend.hpp b/src/core/ggml_extend.hpp index 01d51644..df092479 100644 --- a/src/core/ggml_extend.hpp +++ b/src/core/ggml_extend.hpp @@ -20,14 +20,16 @@ #include #include #include +#include #include #include #include +#include "core/compute_workspace.h" #include "core/ggml_extend_backend.h" #include "core/ggml_graph_cut.h" #include "core/layer_split_partition.h" -#include "core/layer_stream_prefetch.h" +#include "core/runner_cache.h" #include "ggml-alloc.h" #include "ggml-backend.h" #include "ggml.h" @@ -38,7 +40,7 @@ #include "core/rng.hpp" #include "core/tensor_ggml.hpp" #include "core/util.h" -#include "weight_manager.h" +#include "device_residency_manager.h" #define EPS 1e-05f @@ -1772,39 +1774,42 @@ struct GGMLRunnerContext { }; struct GGMLRunner { +private: + std::map logged_compute_bytes_; + size_t logged_segment_count_ = 0; + + sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, size_t direct_bytes); + std::vector memory_requests(const std::vector& sizes, + size_t pending_cache_bytes) const; + bool fits(const std::vector& requests, + const std::vector& params) const; + bool execute_segment(ggml_cgraph* graph, int n_threads); + std::optional> execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function& read_outputs); + protected: typedef std::function get_graph_cb_t; - using GraphCutSegment = sd::ggml_graph_cut::Segment; - using GraphCutPlan = sd::ggml_graph_cut::Plan; + using GraphCutPlan = sd::ggml_graph_cut::Plan; ggml_backend_t runtime_backend = nullptr; ggml_context* params_ctx = nullptr; - ggml_context* cache_ctx = nullptr; - ggml_backend_buffer_t cache_buffer = nullptr; + sd::RunnerCache cache_; + sd::GraphCutTensorCache cut_cache_; + sd::ComputeWorkspace workspace_; + ggml_context* compute_ctx = nullptr; + bool runner_started_ = false; + bool graph_active_ = false; - ggml_context* compute_ctx = nullptr; - ggml_gallocr* compute_allocr = nullptr; - - size_t max_graph_vram_bytes = 0; - bool stream_layers_enabled = false; - bool layer_prefetch_enabled = true; - size_t observed_max_effective_budget_ = 0; - bool graph_cut_layer_split_enabled = false; + size_t max_graph_vram_bytes = 0; + bool graph_cut_layer_split_enabled = false; std::vector graph_cut_layer_split_backend_vram_limits_; std::vector extra_runtime_backends; // borrowed (SDBackendManager-owned) - ggml_backend_sched_t sched = nullptr; // owned - size_t sched_graph_capacity = 0; - ggml_backend_t cpu_fallback_backend = nullptr; // owned, sched requires a trailing CPU backend bool multi_device_eval_callback_warned = false; std::shared_ptr weight_adapter = nullptr; - std::weak_ptr weight_manager; - std::unordered_set kept_compute_param_tensor_set; - std::vector runner_param_tensors; - std::unordered_set runner_param_tensor_set; + std::weak_ptr residency_manager; bool params_tensor_set_dirty_ = true; std::vector one_vec = {1.f}; @@ -1814,7 +1819,6 @@ protected: ggml_tensor* zero_int_tensor = nullptr; std::map backend_tensor_data_map; - std::map cache_tensor_map; // name -> tensor std::vector> debug_tensors; const std::string final_result_name = "ggml_runner_final_result_tensor"; @@ -1876,23 +1880,6 @@ protected: params_tensor_set_dirty_ = true; } - void alloc_cache_ctx() { - ggml_init_params params; - params.mem_size = static_cast(MAX_PARAMS_TENSOR_NUM * ggml_tensor_overhead()); - params.mem_buffer = nullptr; - params.no_alloc = true; - - cache_ctx = ggml_init(params); - GGML_ASSERT(cache_ctx != nullptr); - } - - void free_cache_ctx() { - if (cache_ctx != nullptr) { - ggml_free(cache_ctx); - cache_ctx = nullptr; - } - } - void alloc_compute_ctx() { ggml_init_params params; params.mem_size = static_cast(ggml_tensor_overhead() * MAX_GRAPH_SIZE + ggml_graph_overhead()); @@ -1961,64 +1948,13 @@ protected: return used_params; } - bool prepare_execute_graph_weights(ggml_cgraph* gf, - std::vector& graph_param_tensors, - std::vector& params_to_prepare, - bool keep_compute_params) { - graph_param_tensors = collect_used_param_tensors(gf); - params_to_prepare.clear(); - params_to_prepare.reserve(graph_param_tensors.size()); - for (ggml_tensor* param : graph_param_tensors) { - if (param == nullptr) { - continue; - } - if (keep_compute_params && - kept_compute_param_tensor_set.find(param) != kept_compute_param_tensor_set.end()) { - continue; - } - params_to_prepare.push_back(param); - } - auto manager = weight_manager.lock(); - if (manager == nullptr) { - if (!params_to_prepare.empty()) { - LOG_ERROR("%s weight manager is not set for graph params", get_desc().c_str()); - return false; - } - return true; - } - - if (!manager->prepare_params(params_to_prepare)) { - LOG_ERROR("%s prepare graph weights failed", get_desc().c_str()); - return false; - } - for (ggml_tensor* param : params_to_prepare) { - if (param == nullptr) { - continue; - } - if (runner_param_tensor_set.insert(param).second) { - runner_param_tensors.push_back(param); - } - } - return true; - } - - void free_compute_backend_param_tensors(const std::vector& tensors) { + void evict_compute_backend_param_tensors(const std::vector& tensors) { if (tensors.empty()) { return; } - auto manager = weight_manager.lock(); + auto manager = residency_manager.lock(); if (manager != nullptr) { - manager->release_compute_backend_params(tensors); - } - } - - void free_params_backend_param_tensors(const std::vector& tensors) { - if (tensors.empty()) { - return; - } - auto manager = weight_manager.lock(); - if (manager != nullptr) { - manager->release_params_backend_params(tensors); + manager->evict_compute_backend_params(tensors); } } @@ -2047,6 +1983,9 @@ protected: ggml_cgraph* get_compute_graph(get_graph_cb_t get_graph) { prepare_build_in_tensor_before(); ggml_cgraph* gf = get_graph(); + if (gf == nullptr) { + return nullptr; + } if (ggml_graph_n_nodes(gf) > 0) { auto result = ggml_graph_node(gf, -1); ggml_set_name(result, final_result_name.c_str()); @@ -2056,7 +1995,7 @@ protected: ggml_build_forward_expand(gf, entry.first); } } - for (const auto& entry : cache_tensor_map) { + for (const auto& entry : cache_.outputs()) { if (entry.second != nullptr) { ggml_build_forward_expand(gf, entry.second); } @@ -2080,63 +2019,6 @@ protected: return true; } - // Pass explicit buffer types: synthesized defaults can make CUDA devices - // report supporting each other's buffers and skip a required copy. - bool ensure_sched(ggml_cgraph* gf) { - const size_t required_graph_size = gf != nullptr - ? std::max(1, - (size_t)ggml_graph_n_nodes(gf) + - sd::ggml_graph_cut::leaf_count(gf)) - : 1; - if (sched != nullptr && sched_graph_capacity >= required_graph_size) { - return true; - } - if (sched != nullptr) { - ggml_backend_sched_free(sched); - sched = nullptr; - sched_graph_capacity = 0; - } - std::vector backends; - backends.reserve(extra_runtime_backends.size() + 2); - backends.push_back(runtime_backend); - for (ggml_backend_t backend : extra_runtime_backends) { - backends.push_back(backend); - } - if (cpu_fallback_backend == nullptr && !sd_backend_is_cpu(runtime_backend)) { - cpu_fallback_backend = sd_backend_cpu_init(); - } - if (cpu_fallback_backend != nullptr) { - backends.push_back(cpu_fallback_backend); - } - - std::vector bufts; - bufts.reserve(backends.size()); - ggml_backend_dev_t main_dev = ggml_backend_get_device(runtime_backend); - for (ggml_backend_t backend : backends) { - ggml_backend_buffer_type_t buft = nullptr; - if (backend == cpu_fallback_backend && main_dev != nullptr) { - buft = ggml_backend_dev_host_buffer_type(main_dev); - } - if (buft == nullptr) { - buft = ggml_backend_get_default_buffer_type(backend); - } - bufts.push_back(buft); - } - - sched = ggml_backend_sched_new(backends.data(), - bufts.data(), - (int)backends.size(), - required_graph_size, - /*parallel=*/false, - /*op_offload=*/false); - if (sched == nullptr) { - LOG_ERROR("%s: failed to create backend sched", get_desc().c_str()); - return false; - } - sched_graph_capacity = required_graph_size; - return true; - } - ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const { if (tensor == nullptr || tensor->buffer == nullptr) { return nullptr; @@ -2163,7 +2045,7 @@ protected: // Weightless ops have no scheduler anchor, so pin them to the most recent // weight device. Views must stay unpinned or cross-device copies can be // skipped for their consumers. - void pin_multi_device_nodes(ggml_cgraph* gf) { + void pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph = nullptr) { if (sched == nullptr || gf == nullptr) { return; } @@ -2171,7 +2053,7 @@ protected: const int n_nodes = ggml_graph_n_nodes(gf); for (int i = 0; i < n_nodes; i++) { ggml_tensor* node = ggml_graph_node(gf, i); - auto node_assignment = graph_cut_layer_split_node_assignments_.find(node); + auto node_assignment = graph_cut_layer_split_node_assignments_.find(original_graph == nullptr ? node : ggml_graph_node(original_graph, i)); if (node_assignment != graph_cut_layer_split_node_assignments_.end()) { current = node_assignment->second; } @@ -2197,152 +2079,31 @@ protected: return !extra_runtime_backends.empty(); } - bool graph_requires_backend_fallback(ggml_cgraph* gf) const { - if (gf == nullptr || sd_backend_is_cpu(runtime_backend)) { - return false; - } - const int n_nodes = ggml_graph_n_nodes(gf); - for (int i = 0; i < n_nodes; ++i) { - ggml_tensor* node = ggml_graph_node(gf, i); - if (node != nullptr && !ggml_backend_supports_op(runtime_backend, node)) { - return true; - } - } - return false; + size_t reusable_compute_buffer_bytes() const { + return workspace_.bytes(runtime_backend); } - bool alloc_compute_buffer(ggml_cgraph* gf) { - if (sched != nullptr || is_multi_device() || graph_requires_backend_fallback(gf)) { - // The sched replaces the gallocr. Do NOT ggml_backend_sched_reserve - // the graph here: reserve runs split_graph, which rewires the - // graph's src pointers to sched-internal copy tensors, and the - // later ggml_backend_sched_alloc_graph would split the already - // rewired graph, silently corrupting every cross-backend input. A - // graph must be split at most once; the alloc in execute_graph - // performs the real allocation. - if (compute_allocr != nullptr) { - ggml_gallocr_free(compute_allocr); - compute_allocr = nullptr; - } - return ensure_sched(gf); + size_t retained_runtime_buffer_bytes(ggml_backend_t backend = nullptr) const { + backend = backend == nullptr ? runtime_backend : backend; + size_t bytes = workspace_.bytes(backend); + if (backend == runtime_backend) { + const size_t cache_bytes = cache_.resident_bytes(ggml_backend_get_device(backend)); + bytes = cache_bytes > SIZE_MAX - bytes ? SIZE_MAX : bytes + cache_bytes; + const size_t cut_bytes = cut_cache_.resident_bytes(ggml_backend_get_device(backend)); + bytes = cut_bytes > SIZE_MAX - bytes ? SIZE_MAX : bytes + cut_bytes; } - if (compute_allocr != nullptr) { - return true; - } - compute_allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_backend)); - - if (!ggml_gallocr_reserve(compute_allocr, gf)) { - // failed to allocate the compute buffer - LOG_ERROR("%s: failed to allocate the compute buffer\n", get_desc().c_str()); - free_compute_buffer(); - return false; - } - - // compute the required memory - size_t compute_buffer_size = ggml_gallocr_get_buffer_size(compute_allocr, 0); - LOG_DEBUG("%s compute buffer size: %.2f MB(%s)", - get_desc().c_str(), - compute_buffer_size / 1024.0 / 1024.0, - sd_backend_is_cpu(runtime_backend) ? "RAM" : "VRAM"); - return true; + return bytes; } - void free_cache_buffer() { - if (cache_buffer != nullptr) { - ggml_backend_buffer_free(cache_buffer); - cache_buffer = nullptr; - } - } - - bool copy_cache_tensors_to_cache_buffer(const std::unordered_set* cache_keep_names = nullptr) { - if (cache_tensor_map.empty() && cache_keep_names == nullptr) { - return true; - } - - ggml_context* old_cache_ctx = cache_ctx; - ggml_backend_buffer_t old_cache_buffer = cache_buffer; - cache_ctx = nullptr; - cache_buffer = nullptr; - std::map merged_cache_sources; - if (old_cache_ctx != nullptr) { - for (ggml_tensor* tensor = ggml_get_first_tensor(old_cache_ctx); tensor != nullptr; tensor = ggml_get_next_tensor(old_cache_ctx, tensor)) { - if (cache_keep_names != nullptr && cache_keep_names->find(tensor->name) == cache_keep_names->end()) { - continue; - } - merged_cache_sources[tensor->name] = tensor; + void sync_runtime_residency() { + if (auto manager = residency_manager.lock()) { + manager->update_runtime_residency(reinterpret_cast(this), + runtime_backend, retained_runtime_buffer_bytes()); + for (auto backend : extra_runtime_backends) { + manager->update_runtime_residency(reinterpret_cast(this), + backend, retained_runtime_buffer_bytes(backend)); } } - for (const auto& kv : cache_tensor_map) { - if (cache_keep_names != nullptr && cache_keep_names->find(kv.first) == cache_keep_names->end()) { - continue; - } - merged_cache_sources[kv.first] = kv.second; - } - cache_tensor_map.clear(); - if (merged_cache_sources.empty()) { - if (old_cache_buffer != nullptr) { - ggml_backend_buffer_free(old_cache_buffer); - } - if (old_cache_ctx != nullptr) { - ggml_free(old_cache_ctx); - } - return true; - } - - alloc_cache_ctx(); - std::vector> source_to_cache_tensors; - source_to_cache_tensors.reserve(merged_cache_sources.size()); - for (const auto& kv : merged_cache_sources) { - ggml_tensor* source_tensor = sd::ggml_graph_cut::cache_source_tensor(kv.second); - auto cache_tensor = ggml_dup_tensor(cache_ctx, source_tensor); - ggml_set_name(cache_tensor, kv.first.c_str()); - source_to_cache_tensors.push_back({source_tensor, cache_tensor}); - } - size_t num_tensors = ggml_tensor_num(cache_ctx); - cache_buffer = ggml_backend_alloc_ctx_tensors(cache_ctx, runtime_backend); - GGML_ASSERT(cache_buffer != nullptr); - for (const auto& kv : source_to_cache_tensors) { - ggml_tensor* src = kv.first; - ggml_tensor* dst = kv.second; - ggml_backend_buffer_t src_buf = sd::ggml_graph_cut::tensor_buffer(src); - ggml_backend_buffer_t dst_buf = sd::ggml_graph_cut::tensor_buffer(dst); - if (src_buf == nullptr || dst_buf == nullptr) { - LOG_ERROR("%s cache copy tensor buffer missing: name=%s op=%s src0=%p src0_name=%s src0_buffer=%p src_buffer=%p src_view_src=%p src_view_src_buffer=%p dst_buffer=%p", - get_desc().c_str(), - src && src->name[0] != '\0' ? src->name : "", - src ? ggml_op_name(src->op) : "", - src ? src->src[0] : nullptr, - (src && src->src[0] && src->src[0]->name[0] != '\0') ? src->src[0]->name : "", - (src && src->src[0]) ? sd::ggml_graph_cut::tensor_buffer(src->src[0]) : nullptr, - src ? src->buffer : nullptr, - src ? src->view_src : nullptr, - (src && src->view_src) ? src->view_src->buffer : nullptr, - dst ? dst->buffer : nullptr); - return false; - } - const bool use_staging_copy = src->view_src != nullptr || !ggml_is_contiguous(src) || src->buffer == nullptr; - if (use_staging_copy) { - std::vector host_data(ggml_nbytes(src)); - ggml_backend_tensor_get(src, host_data.data(), 0, host_data.size()); - ggml_backend_tensor_set(dst, host_data.data(), 0, host_data.size()); - } else { - ggml_backend_tensor_copy(src, dst); - } - } - ggml_backend_synchronize(runtime_backend); - size_t cache_buffer_size = ggml_backend_buffer_get_size(cache_buffer); - LOG_DEBUG("%s cache backend buffer size = % 6.2f MB(%s) (%i tensors)", - get_desc().c_str(), - cache_buffer_size / (1024.f * 1024.f), - sd_backend_is_cpu(runtime_backend) ? "RAM" : "VRAM", - num_tensors); - if (old_cache_buffer != nullptr) { - ggml_backend_buffer_free(old_cache_buffer); - } - if (old_cache_ctx != nullptr) { - ggml_free(old_cache_ctx); - } - return true; } template @@ -2372,13 +2133,7 @@ protected: return std::nullopt; } - sd::Tensor result(sd::shape_from_ggml(tensor)); - if (tensor->view_src != nullptr || !ggml_is_contiguous(tensor) || tensor->buffer == nullptr) { - ggml_backend_tensor_get(tensor, result.data(), 0, ggml_nbytes(tensor)); - } else { - ggml_backend_tensor_get(tensor, result.data(), 0, ggml_nbytes(tensor)); - } - return result; + return sd::make_sd_tensor_from_ggml(tensor); } void copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy = true) { @@ -2436,114 +2191,21 @@ protected: } } - bool should_use_graph_cut_segmented_compute(const GraphCutPlan& plan) { - return plan.has_cuts && - plan.valid && - max_graph_vram_bytes > 0 && - plan.segments.size() > 1 && - !sd_backend_is_cpu(runtime_backend) && - !is_multi_device(); - } - - bool can_attempt_graph_cut_segmented_compute() const { - return max_graph_vram_bytes > 0 && - !sd_backend_is_cpu(runtime_backend) && - !is_multi_device(); - } - bool resolve_graph_cut_plan(ggml_cgraph* gf, - GraphCutPlan* plan_out, - size_t* effective_budget_out = nullptr) { + GraphCutPlan* plan_out) { GGML_ASSERT(plan_out != nullptr); GGML_ASSERT(gf != nullptr); - - size_t effective_budget = max_graph_vram_bytes; - size_t free_clamp = SIZE_MAX; - if (stream_layers_enabled && max_graph_vram_bytes > 0 && runtime_backend != nullptr) { - ggml_backend_dev_t dev = ggml_backend_get_device(runtime_backend); - if (dev != nullptr && ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) { - size_t free_vram = 0, total_vram = 0; - ggml_backend_dev_memory(dev, &free_vram, &total_vram); - constexpr size_t safety_margin = 512ull * 1024 * 1024; - free_clamp = (free_vram > safety_margin) ? (free_vram - safety_margin) : 0; - if (free_clamp < effective_budget) { - LOG_DEBUG("%s clamping streaming budget: actual free VRAM %.2f MB < user cap %.2f MB", - get_desc().c_str(), - free_clamp / (1024.0 * 1024.0), - effective_budget / (1024.0 * 1024.0)); - effective_budget = free_clamp; - } - } - } - - bool budget_increased = false; - if (stream_layers_enabled) { - if (effective_budget > observed_max_effective_budget_) { - observed_max_effective_budget_ = effective_budget; - budget_increased = true; - } else { - // Keep the plan cache stable, but never plan above what is free now: - // another model or process can take VRAM after the first measurement. - effective_budget = std::min(observed_max_effective_budget_, free_clamp); - } - } - - if (effective_budget_out != nullptr) { - *effective_budget_out = effective_budget; - } - - // When streaming and the model dwarfs the budget, cap the planner at - // a quarter so it builds smaller merged segments and chunk-K can fit - // alongside. Without streaming the cap only adds dispatch overhead. - size_t planner_budget = effective_budget; - if (stream_layers_enabled) { - size_t total_params_bytes = 0; - for (const ggml_tensor* t : params_tensor_set_) { - if (t != nullptr) { - total_params_bytes += ggml_nbytes(t); - } - } - if (total_params_bytes * 4 > effective_budget * 3) { - planner_budget = effective_budget / 4; - } - } - *plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend, gf, &graph_cut_plan_cache_, - planner_budget, params_tensor_set_, get_desc().c_str()); - if (stream_layers_enabled) { - sd::ggml_graph_cut::annotate_residency(*plan_out, - effective_budget, - layer_prefetch_enabled); - } - if (stream_layers_enabled) { - if (budget_increased) { - LOG_INFO("%s streaming budget = %.2f MB", - get_desc().c_str(), - effective_budget / (1024.0 * 1024.0)); - } else { - LOG_DEBUG("%s streaming budget = %.2f MB", - get_desc().c_str(), - effective_budget / (1024.0 * 1024.0)); - } - } return true; } bool resolve_graph_cut_layer_split_plan(ggml_cgraph* gf, GraphCutPlan* plan_out) { - GGML_ASSERT(plan_out != nullptr); - GGML_ASSERT(gf != nullptr); - *plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend, - gf, - &graph_cut_plan_cache_, - 0, - params_tensor_set_, - get_desc().c_str()); - return true; + return resolve_graph_cut_plan(gf, plan_out); } bool assign_graph_cut_layer_split_backends(ggml_cgraph* gf) { @@ -2561,7 +2223,7 @@ protected: return false; } if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) { - auto manager = weight_manager.lock(); + auto manager = residency_manager.lock(); if (manager == nullptr) { LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str()); return false; @@ -2610,7 +2272,7 @@ protected: } } - auto manager = weight_manager.lock(); + auto manager = residency_manager.lock(); if (manager == nullptr) { LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str()); return false; @@ -2650,429 +2312,68 @@ protected: return true; } - struct PersistentExternalBinding { - ggml_backend_buffer_t buffer = nullptr; - void* data = nullptr; - void* extra = nullptr; - }; - - void snapshot_persistent_externals(const sd::ggml_graph_cut::Plan& plan, - ggml_cgraph* gf, - std::unordered_map& out) { - GGML_ASSERT(gf != nullptr); - out.clear(); - for (const auto& segment : plan.segments) { - for (const auto& input : segment.input_refs) { - if (input.type != GraphCutSegment::INPUT_EXTERNAL) { - continue; - } - ggml_tensor* tensor = sd::ggml_graph_cut::input_tensor(gf, input); - if (tensor == nullptr || tensor->buffer == nullptr) { - continue; - } - PersistentExternalBinding binding; - binding.buffer = tensor->buffer; - binding.data = tensor->data; - binding.extra = tensor->extra; - out[tensor] = binding; - } +public: + bool runner_start() { + if (runner_started_) { + return true; } - } - - void reset_segment_runtime_tensors(const GraphCutSegment& segment, - ggml_cgraph* gf, - const std::unordered_map* persistent_externals = nullptr) { - GGML_ASSERT(gf != nullptr); - - for (const auto& input : segment.input_refs) { - ggml_tensor* input_tensor = sd::ggml_graph_cut::input_tensor(gf, input); - if (input_tensor == nullptr) { - continue; - } - switch (input.type) { - case GraphCutSegment::INPUT_PREVIOUS_CUT: - input_tensor->buffer = nullptr; - input_tensor->data = nullptr; - input_tensor->extra = nullptr; - break; - case GraphCutSegment::INPUT_EXTERNAL: { - if (persistent_externals != nullptr) { - auto it = persistent_externals->find(input_tensor); - if (it != persistent_externals->end()) { - input_tensor->buffer = it->second.buffer; - input_tensor->data = it->second.data; - input_tensor->extra = it->second.extra; - break; - } - } - input_tensor->buffer = nullptr; - input_tensor->data = nullptr; - input_tensor->extra = nullptr; - break; + cache_.clear(); + workspace_.set_extra_backends(extra_runtime_backends); + if (auto manager = residency_manager.lock()) { + manager->set_workspace_reclaimer(reinterpret_cast(this), [this]() { + if (!workspace_.release()) { + return false; } - case GraphCutSegment::INPUT_PARAM: - break; - } - } - - for (int node_idx : segment.internal_node_indices) { - ggml_tensor* node = ggml_graph_node(gf, node_idx); - if (node == nullptr) { - continue; - } - node->buffer = nullptr; - node->data = nullptr; - node->extra = nullptr; - } - } - - bool bind_segment_cached_inputs(ggml_cgraph* gf, const GraphCutSegment& segment) { - GGML_ASSERT(gf != nullptr); - for (const auto& input : segment.input_refs) { - ggml_tensor* input_tensor = sd::ggml_graph_cut::input_tensor(gf, input); - if (input_tensor == nullptr) { - continue; - } - switch (input.type) { - case GraphCutSegment::INPUT_PREVIOUS_CUT: { - ggml_tensor* cache_tensor = get_cache_tensor_by_name(input.display_name); - if (cache_tensor == nullptr) { - LOG_ERROR("%s missing graph cut cache tensor: %s", - get_desc().c_str(), - input.display_name.c_str()); - return false; - } - if (input_tensor->view_src != nullptr) { - input_tensor->view_src = cache_tensor; - input_tensor->buffer = nullptr; - input_tensor->data = cache_tensor->data == nullptr - ? nullptr - : static_cast(static_cast(cache_tensor->data) + input_tensor->view_offs); - input_tensor->extra = cache_tensor->extra; - } else { - input_tensor->buffer = cache_tensor->buffer; - input_tensor->data = cache_tensor->data; - input_tensor->extra = cache_tensor->extra; - } - for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { - input_tensor->src[src_idx] = nullptr; - } - input_tensor->op = GGML_OP_NONE; - break; - } - case GraphCutSegment::INPUT_EXTERNAL: - case GraphCutSegment::INPUT_PARAM: - break; - } + sync_runtime_residency(); + return true; + }); } + runner_started_ = true; return true; } - template - std::optional> execute_graph(ggml_cgraph* gf, - int n_threads, - bool free_compute_buffer, - bool free_compute_params, - bool preserve_backend_tensor_data_map, - bool no_return = false, - const std::unordered_set* cache_keep_names = nullptr, - const std::function& before_compute = {}) { - std::vector graph_param_tensors; - std::vector params_to_prepare; - if (!prepare_execute_graph_weights(gf, graph_param_tensors, params_to_prepare, !free_compute_params)) { - return std::nullopt; + bool runner_started() const { return runner_started_; } + + void runner_end() { + GGML_ASSERT(!graph_active_); + if (!runner_started_) { + return; } - struct GraphWeightDoneGuard { - GraphWeightDoneGuard(GGMLRunner* runner, const std::vector* tensors) - : runner(runner), - tensors(tensors) {} - - GGMLRunner* runner = nullptr; - const std::vector* tensors = nullptr; - bool enabled = true; - - ~GraphWeightDoneGuard() { - if (enabled && runner != nullptr && tensors != nullptr) { - runner->free_compute_backend_param_tensors(*tensors); - } + workspace_.release(); + cache_.clear(); + logged_compute_bytes_.clear(); + logged_segment_count_ = 0; + if (auto manager = residency_manager.lock()) { + manager->clear_prefetched_params(reinterpret_cast(this)); + std::vector tensors; + for (auto tensor = ggml_get_first_tensor(params_ctx); tensor != nullptr; + tensor = ggml_get_next_tensor(params_ctx, tensor)) { + tensors.push_back(tensor); } - - void dismiss() { enabled = false; } - - GraphWeightDoneGuard(const GraphWeightDoneGuard&) = delete; - GraphWeightDoneGuard& operator=(const GraphWeightDoneGuard&) = delete; - }; - GraphWeightDoneGuard graph_weight_done_guard(this, ¶ms_to_prepare); - - if (!alloc_compute_buffer(gf)) { - LOG_ERROR("%s alloc compute buffer failed", get_desc().c_str()); - return std::nullopt; + manager->evict_compute_backend_params(tensors); + manager->remove_runtime_owner(reinterpret_cast(this)); } - struct ComputeBufferGuard { - ComputeBufferGuard(GGMLRunner* runner, bool enabled) - : runner(runner), - enabled(enabled) {} - - GGMLRunner* runner = nullptr; - bool enabled = false; - - ~ComputeBufferGuard() { - if (enabled && runner != nullptr) { - runner->free_compute_buffer(); - } - } - - ComputeBufferGuard(const ComputeBufferGuard&) = delete; - ComputeBufferGuard& operator=(const ComputeBufferGuard&) = delete; - }; - ComputeBufferGuard compute_buffer_guard(this, free_compute_buffer); - - if (sched != nullptr) { - ggml_backend_sched_reset(sched); - pin_multi_device_nodes(gf); // reset clears the pins; re-apply before alloc - if (!ggml_backend_sched_alloc_graph(sched, gf)) { - LOG_ERROR("%s sched alloc compute graph failed", get_desc().c_str()); - return std::nullopt; - } - } else if (!ggml_gallocr_alloc_graph(compute_allocr, gf)) { - LOG_ERROR("%s alloc compute graph failed", get_desc().c_str()); - return std::nullopt; - } - - copy_data_to_backend_tensor(gf, !preserve_backend_tensor_data_map); - if (before_compute) { - before_compute(); - } - if (sd_backend_is_cpu(runtime_backend)) { - sd_backend_cpu_set_n_threads(runtime_backend, n_threads); - } - if (cpu_fallback_backend != nullptr) { - sd_backend_cpu_set_n_threads(cpu_fallback_backend, n_threads); - } - - ggml_status status; - if (sched != nullptr) { - if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) { - LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring", - get_desc().c_str()); - multi_device_eval_callback_warned = true; - } - status = ggml_backend_sched_graph_compute(sched, gf); - if (status == GGML_STATUS_SUCCESS) { - ggml_backend_sched_synchronize(sched); - } - } else { - status = sd_backend_graph_compute_with_eval_callback(runtime_backend, - gf, - sd_get_backend_eval_callback(), - sd_get_backend_eval_callback_data()); - } - if (status != GGML_STATUS_SUCCESS) { - LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status)); - return std::nullopt; - } - - if (!debug_tensors.empty()) { - std::unordered_set debug_graph_tensor_set; - const int n_debug_leafs = sd::ggml_graph_cut::leaf_count(gf); - const int n_debug_nodes = ggml_graph_n_nodes(gf); - debug_graph_tensor_set.reserve(static_cast(n_debug_leafs + n_debug_nodes)); - for (int i = 0; i < n_debug_leafs; ++i) { - debug_graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i)); - } - for (int i = 0; i < n_debug_nodes; ++i) { - debug_graph_tensor_set.insert(ggml_graph_node(gf, i)); - } - - for (const auto& entry : debug_tensors) { - auto tensor = entry.first; - if (tensor == nullptr) { - continue; - } - if (debug_graph_tensor_set.find(tensor) == debug_graph_tensor_set.end()) { - continue; - } - ggml_backend_buffer_t tensor_buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - if (tensor_buf == nullptr) { - LOG_WARN("%s skip debug tensor '%s': tensor buffer not set", - get_desc().c_str(), - entry.second.c_str()); - continue; - } - if (tensor->type != GGML_TYPE_F32) { - LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s", - get_desc().c_str(), - entry.second.c_str(), - ggml_type_name(tensor->type)); - continue; - } - auto debug_tensor = sd::make_sd_tensor_from_ggml(tensor); - print_sd_tensor(debug_tensor, false, entry.second.c_str()); - } - } - - if (!copy_cache_tensors_to_cache_buffer(cache_keep_names)) { - return std::nullopt; - } - auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str()); - std::optional> output; - if (!no_return) { - output = read_graph_tensor(result, "output"); - if (!output.has_value()) { - return std::nullopt; - } - } else { - output = sd::Tensor(); - } - - if (!free_compute_params) { - for (ggml_tensor* param : params_to_prepare) { - if (param == nullptr) { - continue; - } - kept_compute_param_tensor_set.insert(param); - } - graph_weight_done_guard.dismiss(); - } - return output; - } - - template - std::optional> compute_graph_cut_segments(ggml_cgraph* gf, - const GraphCutPlan& plan, - int n_threads, - bool log_residency, - bool no_return = false) { - GGML_ASSERT(gf != nullptr); - - free_compute_buffer(); - free_cache_ctx_and_buffer(); - - sd::LayerStreamPrefetch prefetch(weight_manager.lock(), - reinterpret_cast(this), - gf, - plan, - params_tensor_set_, - stream_layers_enabled && layer_prefetch_enabled); - auto disable_prefetch = [&]() { - if (layer_prefetch_enabled) { - LOG_WARN("%s layer prefetch failed; continuing with synchronous streaming", - get_desc().c_str()); - } - layer_prefetch_enabled = false; - }; - - std::unordered_map persistent_externals; - snapshot_persistent_externals(plan, gf, persistent_externals); - - std::optional> output = sd::Tensor(); - for (size_t seg_idx = 0; seg_idx < plan.segments.size(); ++seg_idx) { - const auto& segment = plan.segments[seg_idx]; - const bool is_last = seg_idx + 1 == plan.segments.size(); - auto future_cut_names = sd::ggml_graph_cut::collect_future_input_names(gf, plan, seg_idx); - if (!prefetch.activate(seg_idx)) { - disable_prefetch(); - } - if (log_residency) { - LOG_DEBUG("%s graph cut executing segment %zu/%zu: %s (residency=%s)", - get_desc().c_str(), - seg_idx + 1, - plan.segments.size(), - segment.group_name.c_str(), - segment.residency == sd::ggml_graph_cut::SegmentResidency::RESIDENT ? "RESIDENT" : "STREAMED"); - } else { - LOG_DEBUG("%s graph cut executing segment %zu/%zu: %s", - get_desc().c_str(), - seg_idx + 1, - plan.segments.size(), - segment.group_name.c_str()); - } - - reset_segment_runtime_tensors(segment, gf, &persistent_externals); - if (!bind_segment_cached_inputs(gf, segment)) { - free_cache_ctx_and_buffer(); - free_compute_buffer(); - free_compute_ctx(); - return std::nullopt; - } - - if (!is_last) { - for (size_t output_idx = 0; output_idx < segment.output_node_indices.size(); ++output_idx) { - ggml_tensor* output_tensor = sd::ggml_graph_cut::output_tensor(gf, segment, output_idx); - if (output_tensor != nullptr && - sd::ggml_graph_cut::is_graph_cut_tensor(output_tensor) && - future_cut_names.find(output_tensor->name) != future_cut_names.end()) { - cache(output_tensor->name, output_tensor); - } - } - } - - ggml_context* segment_graph_ctx = nullptr; - ggml_cgraph* segment_graph = sd::ggml_graph_cut::build_segment_graph(gf, segment, &segment_graph_ctx); - const bool keep_segment_params = segment.residency == sd::ggml_graph_cut::SegmentResidency::RESIDENT; - std::function before_compute; - if (prefetch.enabled()) { - before_compute = [&, seg_idx]() { - if (!prefetch.enqueue_next(seg_idx)) { - disable_prefetch(); - } - }; - } - auto segment_output = execute_graph(segment_graph, - n_threads, - true, - !keep_segment_params, - true, - !is_last || no_return, - &future_cut_names, - before_compute); - ggml_free(segment_graph_ctx); - if (!segment_output.has_value()) { - free_cache_ctx_and_buffer(); - free_compute_buffer(); - free_compute_ctx(); - return std::nullopt; - } - output = std::move(segment_output); - } - - backend_tensor_data_map.clear(); - free_cache_ctx_and_buffer(); - free_compute_ctx(); - return output; - } - -public: - void runner_done() { - free_compute_buffer(); - std::vector tensors_to_release = std::move(this->runner_param_tensors); - this->runner_param_tensors.clear(); - runner_param_tensor_set.clear(); - kept_compute_param_tensor_set.clear(); - free_compute_backend_param_tensors(tensors_to_release); - free_params_backend_param_tensors(tensors_to_release); + runner_started_ = false; } public: virtual std::string get_desc() = 0; GGMLRunner(ggml_backend_t backend, - std::shared_ptr manager = nullptr) + std::shared_ptr manager = nullptr) : runtime_backend(backend), - weight_manager(manager) { + cache_(backend), + cut_cache_(backend), + workspace_(backend), + residency_manager(manager) { GGML_ASSERT(runtime_backend != nullptr); alloc_params_ctx(); } virtual ~GGMLRunner() { - free_compute_buffer(); - free_params_ctx(); + runner_end(); free_compute_ctx(); - free_cache_ctx_and_buffer(); - if (cpu_fallback_backend != nullptr) { - ggml_backend_free(cpu_fallback_backend); - cpu_fallback_backend = nullptr; - } + free_params_ctx(); } virtual GGMLRunnerContext get_context() { @@ -3104,20 +2405,8 @@ public: public: void free_cache_ctx_and_buffer() { - free_cache_buffer(); - free_cache_ctx(); - } - - void free_compute_buffer() { - if (compute_allocr != nullptr) { - ggml_gallocr_free(compute_allocr); - compute_allocr = nullptr; - } - if (sched != nullptr) { - ggml_backend_sched_free(sched); - sched = nullptr; - sched_graph_capacity = 0; - } + cache_.clear(); + sync_runtime_residency(); } // do copy after alloc graph @@ -3172,75 +2461,70 @@ public: if (tensor != nullptr && tensor->view_src != nullptr) { tensor = ggml_cont(compute_ctx, tensor); } - cache_tensor_map[name] = tensor; + if (tensor != nullptr) { + ggml_set_output(tensor); + } + cache_.stage(name, tensor); } ggml_tensor* get_cache_tensor_by_name(const std::string& name) { - if (cache_ctx == nullptr) { - return nullptr; - } - return ggml_get_tensor(cache_ctx, name.c_str()); + return cache_.get(name); } template std::optional> compute(get_graph_cb_t get_graph, int n_threads, - bool auto_free = true, - bool free_compute_buffer = true, - bool free_compute_params = true, - bool no_return = false) { - struct RunnerDoneGuard { - RunnerDoneGuard(GGMLRunner* runner, bool enabled) - : runner(runner), - enabled(enabled) {} - - ~RunnerDoneGuard() { - if (enabled && runner != nullptr) { - runner->runner_done(); + bool auto_runner_end = true, + bool no_return = false, + const std::function& read_outputs = {}) { + if (graph_active_) { + LOG_ERROR("%s does not support reentrant graph execution", get_desc().c_str()); + return std::nullopt; + } + if (!runner_start()) { + runner_end(); + return std::nullopt; + } + struct RunnerEndGuard { + GGMLRunner& runner; + bool enabled; + ~RunnerEndGuard() { + if (enabled) { + runner.runner_end(); } } + } runner_guard{*this, auto_runner_end}; + graph_active_ = true; + bool success = false; + struct GraphEndGuard { + GGMLRunner& runner; + const bool& success; + ~GraphEndGuard() { + runner.workspace_.segment_end(); + runner.cache_.graph_end(false); + runner.cut_cache_.clear(); + runner.free_compute_ctx(); + runner.graph_active_ = false; + if (!success) { + runner.workspace_.release(); + } + runner.sync_runtime_residency(); + } + } graph_guard{*this, success}; - RunnerDoneGuard(const RunnerDoneGuard&) = delete; - RunnerDoneGuard& operator=(const RunnerDoneGuard&) = delete; - - GGMLRunner* runner = nullptr; - bool enabled = false; - }; - RunnerDoneGuard runner_done_guard(this, auto_free); - - ggml_cgraph* gf = nullptr; - if (!prepare_compute_graph(get_graph, &gf)) { + ggml_cgraph* graph = nullptr; + if (!prepare_compute_graph(get_graph, &graph)) { return std::nullopt; } - GGML_ASSERT(gf != nullptr); rebuild_params_tensor_set(); - - if (!assign_graph_cut_layer_split_backends(gf)) { - free_compute_ctx(); - return std::nullopt; + static_assert(std::is_same::value, + "GGMLRunner currently supports float graph outputs only"); + auto output = execute_graph(graph, n_threads, no_return, read_outputs); + success = output.has_value(); + if (success) { + cache_.graph_end(true); } - - if (can_attempt_graph_cut_segmented_compute()) { - GraphCutPlan plan; - if (!resolve_graph_cut_plan(gf, &plan)) { - free_compute_ctx(); - return std::nullopt; - } - if (should_use_graph_cut_segmented_compute(plan)) { - return compute_graph_cut_segments(gf, - plan, - n_threads, - stream_layers_enabled, - no_return); - } - } - return execute_graph(gf, - n_threads, - free_compute_buffer, - free_compute_params, - false, - no_return, - nullptr); + return output; } void set_flash_attention_enabled(bool enabled) { @@ -3264,19 +2548,6 @@ public: max_graph_vram_bytes = max_vram_bytes; } - void set_stream_layers_enabled(bool enabled) { - if (enabled && is_multi_device()) { - LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring", - get_desc().c_str()); - return; - } - stream_layers_enabled = enabled; - } - - void set_layer_prefetch_enabled(bool enabled) { - layer_prefetch_enabled = enabled; - } - void set_graph_cut_layer_split_enabled(bool enabled) { graph_cut_layer_split_enabled = enabled; if (!enabled) { @@ -3304,14 +2575,10 @@ public: extra_runtime_backends.push_back(backend); } } + workspace_.set_extra_backends(extra_runtime_backends); graph_cut_layer_split_assignments_.clear(); graph_cut_layer_split_node_assignments_.clear(); graph_cut_layer_split_primary_notice_logged_ = false; - if (is_multi_device() && stream_layers_enabled) { - LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring", - get_desc().c_str()); - stream_layers_enabled = false; - } } }; diff --git a/src/core/ggml_graph_cut.cpp b/src/core/ggml_graph_cut.cpp index dec74716..e6fa43df 100644 --- a/src/core/ggml_graph_cut.cpp +++ b/src/core/ggml_graph_cut.cpp @@ -2,6 +2,7 @@ #include #include +#include #include #include #include @@ -67,25 +68,6 @@ namespace sd::ggml_graph_cut { return -1; } - static Plan::InputShape input_shape(const ggml_tensor* tensor) { - Plan::InputShape shape; - if (tensor == nullptr) { - return shape; - } - shape.type = tensor->type; - for (int i = 0; i < GGML_MAX_DIMS; ++i) { - shape.ne[static_cast(i)] = tensor->ne[i]; - } - return shape; - } - - static size_t graph_cut_segment_vram_bytes(const Segment& segment) { - return segment.compute_buffer_size + - segment.input_param_bytes + - segment.input_previous_cut_bytes + - segment.output_bytes; - } - static std::string lower_ascii_copy(std::string value) { std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) { return static_cast(std::tolower(c)); @@ -291,55 +273,6 @@ namespace sd::ggml_graph_cut { return max_vram_bytes_to_gib(resolve_auto_max_vram_bytes(-max_vram, backend)); } - static bool is_segment_output_needed_after(const Plan& plan, - size_t end_segment_index, - int output_node_index) { - if (end_segment_index + 1 >= plan.segments.size()) { - return false; - } - for (size_t seg_idx = end_segment_index + 1; seg_idx < plan.segments.size(); ++seg_idx) { - const auto& segment = plan.segments[seg_idx]; - for (const auto& input_ref : segment.input_refs) { - if (input_ref.type == Segment::INPUT_PREVIOUS_CUT && - input_ref.node_index == output_node_index) { - return true; - } - } - } - return false; - } - - static Segment make_segment_seed(const Plan& plan, - size_t start_segment_index, - size_t end_segment_index) { - GGML_ASSERT(start_segment_index < plan.segments.size()); - GGML_ASSERT(end_segment_index < plan.segments.size()); - GGML_ASSERT(start_segment_index <= end_segment_index); - - Segment seed; - const auto& start_segment = plan.segments[start_segment_index]; - const auto& target_segment = plan.segments[end_segment_index]; - std::unordered_set seen_output_node_indices; - for (size_t seg_idx = start_segment_index; seg_idx <= end_segment_index; ++seg_idx) { - const bool is_boundary_segment = seg_idx == end_segment_index; - for (int output_node_index : plan.segments[seg_idx].output_node_indices) { - if ((is_boundary_segment || - is_segment_output_needed_after(plan, end_segment_index, output_node_index)) && - seen_output_node_indices.insert(output_node_index).second) { - seed.output_node_indices.push_back(output_node_index); - } - } - } - if (start_segment_index == end_segment_index) { - seed.group_name = target_segment.group_name; - } else { - seed.group_name = sd_format("%s..%s", - start_segment.group_name.c_str(), - target_segment.group_name.c_str()); - } - return seed; - } - static void build_segment(ggml_cgraph* gf, Plan& plan, Segment& segment, @@ -416,31 +349,7 @@ namespace sd::ggml_graph_cut { } return a.display_name < b.display_name; }); - segment.input_refs = input_refs; - for (const auto& input : input_refs) { - ggml_tensor* current_input = input_tensor(gf, input); - size_t tensor_bytes = current_input == nullptr - ? 0 - : (input.type == Segment::INPUT_PREVIOUS_CUT - ? cache_tensor_bytes(current_input) - : ggml_nbytes(current_input)); - switch (input.type) { - case Segment::INPUT_PREVIOUS_CUT: - segment.input_previous_cut_bytes += tensor_bytes; - break; - case Segment::INPUT_PARAM: - segment.input_param_bytes += tensor_bytes; - break; - case Segment::INPUT_EXTERNAL: - default: - segment.input_external_bytes += tensor_bytes; - break; - } - } - for (int output_node_index : segment.output_node_indices) { - ggml_tensor* output = ggml_graph_node(gf, output_node_index); - segment.output_bytes += cache_tensor_bytes(output); - } + segment.input_refs = input_refs; segment.compute_buffer_size = measure_segment_compute_buffer(backend, gf, segment, log_desc); for (int output_node_index : segment.output_node_indices) { @@ -449,6 +358,70 @@ namespace sd::ggml_graph_cut { plan.segments.push_back(std::move(segment)); } + static bool validate_plan(ggml_cgraph* gf, + const Plan& plan, + std::string* validation_error) { + auto fail = [&](const std::string& reason) { + if (validation_error != nullptr) { + *validation_error = reason; + } + return false; + }; + if (!plan.has_cuts) { + return true; + } + if (plan.segments.size() <= 1) { + return fail("fewer than two segments"); + } + const int n_nodes = ggml_graph_n_nodes(gf); + std::unordered_set completed_outputs; + for (size_t segment_index = 0; segment_index < plan.segments.size(); ++segment_index) { + const Segment& segment = plan.segments[segment_index]; + const std::string segment_label = "segment " + std::to_string(segment_index) + + " ('" + segment.group_name + "')"; + if (segment.internal_node_indices.empty() || segment.output_node_indices.empty()) { + return fail(segment_label + " has no internal nodes or outputs"); + } + for (const Segment::InputRef& input : segment.input_refs) { + if (input.type == Segment::INPUT_PREVIOUS_CUT) { + if (input.node_index < 0 || input.node_index >= n_nodes || + completed_outputs.find(input.node_index) == completed_outputs.end()) { + return fail(segment_label + " references an unavailable cut node " + + std::to_string(input.node_index)); + } + } else if (input.leaf_index < 0 || input.leaf_index >= gf->n_leafs) { + return fail(segment_label + " references an invalid leaf " + + std::to_string(input.leaf_index)); + } + } + std::unordered_set segment_nodes; + segment_nodes.reserve(segment.internal_node_indices.size()); + for (int node_index : segment.internal_node_indices) { + if (node_index < 0 || node_index >= n_nodes) { + return fail(segment_label + " contains an invalid node " + + std::to_string(node_index)); + } + if (!segment_nodes.insert(node_index).second) { + return fail(segment_label + " contains duplicate node " + + std::to_string(node_index)); + } + } + for (int output_index : segment.output_node_indices) { + if (output_index < 0 || output_index >= n_nodes || + segment_nodes.find(output_index) == segment_nodes.end()) { + return fail(segment_label + " has an output outside its node set: " + + std::to_string(output_index)); + } + if (completed_outputs.find(output_index) != completed_outputs.end()) { + return fail(segment_label + " repeats output node " + + std::to_string(output_index)); + } + completed_outputs.insert(output_index); + } + } + return true; + } + bool is_graph_cut_tensor(const ggml_tensor* tensor) { if (tensor == nullptr || tensor->name[0] == '\0') { return false; @@ -509,26 +482,80 @@ namespace sd::ggml_graph_cut { return ggml_nbytes(cache_src); } - bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan) { - GGML_ASSERT(gf != nullptr); - if (ggml_graph_n_nodes(gf) != plan.n_nodes || gf->n_leafs != plan.n_leafs) { - return false; - } - for (const auto& input_shape_ref : plan.input_shapes) { - if (input_shape_ref.leaf_index < 0 || input_shape_ref.leaf_index >= gf->n_leafs) { - return false; + std::vector graph_layout(ggml_cgraph* graph, bool include_bindings) { + std::vector tensors; + std::unordered_map indices; + auto add = [&](const ggml_tensor* tensor) { + if (tensor != nullptr && indices.emplace(tensor, tensors.size() + 1).second) { + tensors.push_back(tensor); } - ggml_tensor* leaf = gf->leafs[input_shape_ref.leaf_index]; - if (leaf == nullptr || input_shape_ref.type != leaf->type) { - return false; + }; + for (int i = 0; i < graph->n_leafs; ++i) { + add(graph->leafs[i]); + } + for (int i = 0; i < graph->n_nodes; ++i) { + add(graph->nodes[i]); + } + for (size_t i = 0; i < tensors.size(); ++i) { + add(tensors[i]->view_src); + for (auto source : tensors[i]->src) { + add(source); + } + } + std::vector signature; + signature.reserve(tensors.size() * 24); + signature.push_back(graph->n_nodes); + signature.push_back(graph->n_leafs); + for (int i = 0; i < graph->n_leafs; ++i) { + signature.push_back(indices.at(graph->leafs[i])); + } + for (int i = 0; i < graph->n_nodes; ++i) { + signature.push_back(indices.at(graph->nodes[i])); + } + for (auto tensor : tensors) { + signature.push_back(tensor->op); + signature.push_back(tensor->type); + signature.push_back(tensor->flags); + signature.push_back(tensor->view_offs); + if (include_bindings) { + signature.push_back(tensor->data != nullptr); + auto buffer = tensor_buffer(tensor); + signature.push_back(reinterpret_cast(buffer == nullptr ? nullptr : ggml_backend_buffer_get_type(buffer))); } for (int d = 0; d < GGML_MAX_DIMS; ++d) { - if (input_shape_ref.ne[static_cast(d)] != leaf->ne[d]) { - return false; - } + signature.push_back(tensor->ne[d]); + signature.push_back(tensor->nb[d]); + } + signature.push_back(tensor->view_src == nullptr ? 0 : indices.at(tensor->view_src)); + for (auto source : tensor->src) { + signature.push_back(source == nullptr ? 0 : indices.at(source)); + } + for (int value : tensor->op_params) { + signature.push_back(static_cast(value)); } } - return true; + return signature; + } + + bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan) { + GGML_ASSERT(gf != nullptr); + if (plan.leaf_names.size() != static_cast(gf->n_leafs) || + plan.layout != graph_layout(gf, false)) { + return false; + } + for (int i = 0; i < gf->n_leafs; ++i) { + if (plan.leaf_names[i] != gf->leafs[i]->name) { + return false; + } + } + std::vector> cut_markers; + for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { + auto node = ggml_graph_node(gf, i); + if (is_graph_cut_tensor(node)) { + cut_markers.emplace_back(i, node->name); + } + } + return cut_markers == plan.cut_markers; } ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index) { @@ -578,26 +605,6 @@ namespace sd::ggml_graph_cut { return tensors; } - std::unordered_set collect_future_input_names(ggml_cgraph* gf, - const Plan& plan, - size_t current_segment_index) { - GGML_ASSERT(gf != nullptr); - std::unordered_set future_input_names; - for (size_t seg_idx = current_segment_index + 1; seg_idx < plan.segments.size(); ++seg_idx) { - const auto& segment = plan.segments[seg_idx]; - for (const auto& input_ref : segment.input_refs) { - if (input_ref.type != Segment::INPUT_PREVIOUS_CUT) { - continue; - } - ggml_tensor* current_input = input_tensor(gf, input_ref); - if (current_input != nullptr && current_input->name[0] != '\0') { - future_input_names.insert(current_input->name); - } - } - } - return future_input_names; - } - ggml_cgraph* build_segment_graph(ggml_cgraph* gf, const Segment& segment, ggml_context** graph_ctx_out) { @@ -662,6 +669,10 @@ namespace sd::ggml_graph_cut { continue; } ggml_set_output(output); + if (output->view_src != nullptr) { + // A consumed output view does not keep its storage alive in gallocr. + ggml_set_output(output->view_src); + } } for (int node_idx : segment.internal_node_indices) { ggml_graph_add_node(segment_graph, ggml_graph_node(gf, node_idx)); @@ -716,6 +727,10 @@ namespace sd::ggml_graph_cut { if (output != nullptr && saved_output_flags.find(output) == saved_output_flags.end()) { saved_output_flags[output] = output->flags; } + if (output != nullptr && output->view_src != nullptr && + saved_output_flags.find(output->view_src) == saved_output_flags.end()) { + saved_output_flags[output->view_src] = output->view_src->flags; + } } ggml_context* graph_ctx = nullptr; @@ -744,6 +759,46 @@ namespace sd::ggml_graph_cut { return buffer_size; } + static size_t measure_graph_compute_buffer( + ggml_backend_t backend, + ggml_cgraph* gf, + const std::unordered_set& params_tensor_set) { + struct TensorRuntimeBinding { + ggml_backend_buffer_t buffer = nullptr; + void* data = nullptr; + void* extra = nullptr; + }; + std::unordered_map saved_bindings; + auto mark_external = [&](ggml_tensor* tensor) { + if (tensor == nullptr || saved_bindings.find(tensor) != saved_bindings.end()) { + return; + } + saved_bindings[tensor] = {tensor->buffer, tensor->data, tensor->extra}; + tensor->data = reinterpret_cast(static_cast(1)); + }; + for (int i = 0; i < leaf_count(gf); ++i) { + ggml_tensor* leaf = leaf_tensor(gf, i); + if (!is_params_tensor(params_tensor_set, leaf)) { + continue; + } + mark_external(leaf); + mark_external(leaf->view_src); + } + + ggml_gallocr_t allocr = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + size_t sizes[1] = {0}; + ggml_gallocr_reserve_n_size(allocr, gf, nullptr, nullptr, sizes); + ggml_gallocr_free(allocr); + + for (const auto& kv : saved_bindings) { + kv.first->buffer = kv.second.buffer; + kv.first->data = kv.second.data; + kv.first->extra = kv.second.extra; + } + return sizes[0]; + } + Plan build_plan(ggml_backend_t backend, ggml_cgraph* gf, const std::unordered_set& params_tensor_set, @@ -756,24 +811,22 @@ namespace sd::ggml_graph_cut { if (n_nodes <= 0) { return plan; } - plan.n_nodes = n_nodes; - plan.n_leafs = gf->n_leafs; + plan.layout = graph_layout(gf, false); for (int i = 0; i < gf->n_leafs; ++i) { - ggml_tensor* leaf = gf->leafs[i]; - if (is_params_tensor(params_tensor_set, leaf)) { - continue; - } - auto shape = input_shape(leaf); - shape.leaf_index = i; - plan.input_shapes.push_back(shape); + plan.leaf_names.emplace_back(gf->leafs[i]->name); } + plan.compute_buffer_size = + measure_graph_compute_buffer(backend, gf, params_tensor_set); std::unordered_map producer_index; producer_index.reserve(static_cast(n_nodes)); for (int i = 0; i < n_nodes; ++i) { - producer_index[ggml_graph_node(gf, i)] = i; + ggml_tensor* node = ggml_graph_node(gf, i); + producer_index[node] = i; + if (is_graph_cut_tensor(node)) { + plan.cut_markers.push_back({i, node->name}); + } } - std::vector grouped_segments; std::unordered_map group_to_segment; for (int i = 0; i < n_nodes; ++i) { @@ -824,11 +877,24 @@ namespace sd::ggml_graph_cut { if (final_output_index < 0) { final_output_index = n_nodes - 1; } - ggml_tensor* final_output = final_output_index >= 0 ? ggml_graph_node(gf, final_output_index) : nullptr; - if (final_output != nullptr && available_cut_output_node_indices.find(final_output_index) == available_cut_output_node_indices.end()) { - Segment final_segment; - final_segment.group_name = "ggml_runner.final"; + Segment final_segment; + final_segment.group_name = "ggml_runner.final"; + if (final_output_index >= 0 && + available_cut_output_node_indices.find(final_output_index) == + available_cut_output_node_indices.end()) { final_segment.output_node_indices.push_back(final_output_index); + } + for (int i = 0; i < n_nodes; ++i) { + ggml_tensor* node = ggml_graph_node(gf, i); + if (i == final_output_index || node == nullptr || + (node->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 || + available_cut_output_node_indices.find(i) != + available_cut_output_node_indices.end()) { + continue; + } + final_segment.output_node_indices.push_back(i); + } + if (!final_segment.output_node_indices.empty()) { build_segment(gf, plan, final_segment, @@ -839,223 +905,53 @@ namespace sd::ggml_graph_cut { log_desc); } - return plan; - } - - Plan apply_max_vram_budget(ggml_cgraph* gf, - const Plan& base_plan, - size_t max_graph_vram_bytes, - ggml_backend_t backend, - const std::unordered_set& params_tensor_set, - const char* log_desc) { - GGML_ASSERT(backend != nullptr); - GGML_ASSERT(gf != nullptr); - int64_t t_budget_begin = ggml_time_ms(); - if (max_graph_vram_bytes == 0 || !base_plan.has_cuts || base_plan.segments.size() <= 1) { - return base_plan; - } - - const int n_nodes = ggml_graph_n_nodes(gf); - std::unordered_map producer_index; - producer_index.reserve(static_cast(n_nodes)); - for (int i = 0; i < n_nodes; ++i) { - producer_index[ggml_graph_node(gf, i)] = i; - } - - Plan merged_plan; - merged_plan.available = true; - merged_plan.has_cuts = base_plan.has_cuts; - merged_plan.valid = base_plan.valid; - merged_plan.n_nodes = base_plan.n_nodes; - merged_plan.n_leafs = base_plan.n_leafs; - - std::unordered_set available_cut_output_node_indices; - available_cut_output_node_indices.reserve(static_cast(n_nodes)); - - size_t start_segment_index = 0; - while (start_segment_index < base_plan.segments.size()) { - Plan single_plan; - auto single_available_cut_output_node_indices = available_cut_output_node_indices; - auto single_seed = make_segment_seed(base_plan, - start_segment_index, - start_segment_index); - build_segment(gf, - single_plan, - single_seed, - producer_index, - single_available_cut_output_node_indices, - backend, - params_tensor_set, - log_desc); - GGML_ASSERT(!single_plan.segments.empty()); - - size_t best_end_segment_index = start_segment_index; - bool can_merge_next_segment = graph_cut_segment_vram_bytes(single_plan.segments.back()) <= max_graph_vram_bytes; - - while (can_merge_next_segment && best_end_segment_index + 1 < base_plan.segments.size()) { - const size_t next_end_segment_index = best_end_segment_index + 1; - Plan candidate_plan; - auto candidate_available_cut_output_node_indices = available_cut_output_node_indices; - auto candidate_seed = make_segment_seed(base_plan, - start_segment_index, - next_end_segment_index); - build_segment(gf, - candidate_plan, - candidate_seed, - producer_index, - candidate_available_cut_output_node_indices, - backend, - params_tensor_set, - log_desc); - GGML_ASSERT(!candidate_plan.segments.empty()); - - const auto& candidate_segment = candidate_plan.segments.back(); - const size_t candidate_bytes = graph_cut_segment_vram_bytes(candidate_segment); - if (candidate_bytes > max_graph_vram_bytes) { - break; + std::unordered_set future_cut_names; + for (auto segment = plan.segments.rbegin(); segment != plan.segments.rend(); ++segment) { + segment->future_cut_names = future_cut_names; + segment->live_cut_names = future_cut_names; + for (const auto& input : segment->input_refs) { + if (input.type != Segment::INPUT_PREVIOUS_CUT) { + continue; } - - best_end_segment_index = next_end_segment_index; + segment->live_cut_names.insert(input.display_name); + future_cut_names.insert(input.display_name); } - - auto best_seed = make_segment_seed(base_plan, - start_segment_index, - best_end_segment_index); - build_segment(gf, - merged_plan, - best_seed, - producer_index, - available_cut_output_node_indices, - backend, - params_tensor_set, - log_desc); - start_segment_index = best_end_segment_index + 1; } - if (log_desc != nullptr && merged_plan.segments.size() != base_plan.segments.size()) { - LOG_INFO("%s graph cut max_vram=%.2f MB merged %zu segments -> %zu segments", + std::string plan_validation_error; + plan.valid = validate_plan(gf, plan, &plan_validation_error); + if (!plan.valid && log_desc != nullptr) { + LOG_WARN("%s graph cut plan validation failed (%s); using monolithic execution", log_desc, - max_graph_vram_bytes / 1024.0 / 1024.0, - base_plan.segments.size(), - merged_plan.segments.size()); + plan_validation_error.c_str()); } - if (log_desc != nullptr) { - LOG_DEBUG("%s graph cut max_vram budget merge took %lld ms", - log_desc, - ggml_time_ms() - t_budget_begin); - } - - return merged_plan; + return plan; } Plan resolve_plan(ggml_backend_t backend, ggml_cgraph* gf, PlanCache* cache, - size_t max_graph_vram_bytes, const std::unordered_set& params_tensor_set, const char* log_desc) { GGML_ASSERT(backend != nullptr); GGML_ASSERT(gf != nullptr); GGML_ASSERT(cache != nullptr); - int64_t t_prepare_begin = ggml_time_ms(); - Plan base_plan; - int64_t t_plan_begin = ggml_time_ms(); - if (cache->graph_cut_plan.available && plan_matches_graph(gf, cache->graph_cut_plan)) { - base_plan = cache->graph_cut_plan; - } else { - base_plan = build_plan(backend, gf, params_tensor_set, log_desc); - cache->graph_cut_plan = base_plan; - cache->graph_cut_plan.available = true; - cache->budgeted_graph_cut_plan.available = false; - if (log_desc != nullptr) { - LOG_INFO("%s build cached graph cut plan done (taking %lld ms)", log_desc, ggml_time_ms() - t_plan_begin); - } + if (cache->graph_cut_plan.available && + plan_matches_graph(gf, cache->graph_cut_plan)) { + return cache->graph_cut_plan; } - Plan resolved_plan = base_plan; - if (max_graph_vram_bytes > 0 && base_plan.has_cuts) { - if (cache->budgeted_graph_cut_plan.available && - cache->budgeted_graph_cut_plan_max_vram_bytes == max_graph_vram_bytes && - plan_matches_graph(gf, cache->budgeted_graph_cut_plan)) { - resolved_plan = cache->budgeted_graph_cut_plan; - } else { - resolved_plan = apply_max_vram_budget(gf, - base_plan, - max_graph_vram_bytes, - backend, - params_tensor_set, - log_desc); - cache->budgeted_graph_cut_plan = resolved_plan; - cache->budgeted_graph_cut_plan.available = true; - cache->budgeted_graph_cut_plan_max_vram_bytes = max_graph_vram_bytes; - } - } - return resolved_plan; - } - - void annotate_residency(Plan& plan, - size_t max_graph_vram_bytes, - bool prefetch_enabled) { - // Cached plans may be reused with a smaller live budget. - for (auto& seg : plan.segments) { - seg.residency = SegmentResidency::STREAMED; - } - if (max_graph_vram_bytes == 0 || plan.segments.size() < 2) { - return; - } - - bool any_param_bearing = false; - for (const auto& seg : plan.segments) { - if (seg.input_param_bytes > 0) { - any_param_bearing = true; - break; - } - } - if (!any_param_bearing) { - return; - } - - // Leave room for the largest active streamed segment. - size_t worst_streamed_footprint = 0; - size_t prefetch_headroom = 0; - for (const auto& seg : plan.segments) { - const size_t seg_footprint = seg.input_param_bytes + - seg.compute_buffer_size + - seg.output_bytes + - seg.input_previous_cut_bytes + - seg.input_external_bytes; - if (seg_footprint > worst_streamed_footprint) { - worst_streamed_footprint = seg_footprint; - } - prefetch_headroom = std::max(prefetch_headroom, seg.input_param_bytes); - } - constexpr size_t safety = 512ull * 1024 * 1024; - if (worst_streamed_footprint > SIZE_MAX - safety) { - return; - } - size_t reserved = safety + worst_streamed_footprint; - if (prefetch_enabled) { - if (prefetch_headroom > SIZE_MAX - reserved) { - return; - } - reserved += prefetch_headroom; - } - - if (max_graph_vram_bytes <= reserved) { - return; - } - const size_t available = max_graph_vram_bytes - reserved; - - size_t cumulative = 0; - for (auto& seg : plan.segments) { - if (cumulative + seg.input_param_bytes > available) { - break; - } - seg.residency = SegmentResidency::RESIDENT; - cumulative += seg.input_param_bytes; + int64_t t_plan_begin = ggml_time_ms(); + Plan plan = build_plan(backend, gf, params_tensor_set, log_desc); + cache->graph_cut_plan = plan; + if (log_desc != nullptr) { + LOG_INFO("%s build cached graph cut plan done (taking %lld ms)", + log_desc, + ggml_time_ms() - t_plan_begin); } + return plan; } } // namespace sd::ggml_graph_cut diff --git a/src/core/ggml_graph_cut.h b/src/core/ggml_graph_cut.h index c2594cee..89605a99 100644 --- a/src/core/ggml_graph_cut.h +++ b/src/core/ggml_graph_cut.h @@ -6,19 +6,13 @@ #include #include #include +#include #include #include "ggml-backend.h" #include "ggml.h" namespace sd::ggml_graph_cut { - - // Streaming residency for a segment's params. - enum class SegmentResidency : uint8_t { - STREAMED = 0, - RESIDENT = 1, - }; - struct Segment { enum InputType { INPUT_EXTERNAL = 0, @@ -33,38 +27,28 @@ namespace sd::ggml_graph_cut { int node_index = -1; }; - size_t compute_buffer_size = 0; - size_t output_bytes = 0; - size_t input_external_bytes = 0; - size_t input_previous_cut_bytes = 0; - size_t input_param_bytes = 0; + size_t compute_buffer_size = 0; std::string group_name; std::vector internal_node_indices; std::vector output_node_indices; std::vector input_refs; - SegmentResidency residency = SegmentResidency::STREAMED; + std::unordered_set future_cut_names; + std::unordered_set live_cut_names; }; struct Plan { - struct InputShape { - int leaf_index = -1; - ggml_type type = GGML_TYPE_COUNT; - std::array ne = {0, 0, 0, 0}; - }; - - bool available = false; - bool has_cuts = false; - bool valid = true; - int n_nodes = 0; - int n_leafs = 0; - std::vector input_shapes; + bool available = false; + bool has_cuts = false; + bool valid = true; + size_t compute_buffer_size = 0; + std::vector layout; + std::vector leaf_names; + std::vector> cut_markers; std::vector segments; }; struct PlanCache { Plan graph_cut_plan; - Plan budgeted_graph_cut_plan; - size_t budgeted_graph_cut_plan_max_vram_bytes = 0; }; static constexpr const char* GGML_RUNNER_CUT_PREFIX = "ggml_runner_cut:"; @@ -89,13 +73,12 @@ namespace sd::ggml_graph_cut { ggml_backend_buffer_t tensor_buffer(const ggml_tensor* tensor); ggml_tensor* cache_source_tensor(ggml_tensor* tensor); size_t cache_tensor_bytes(const ggml_tensor* tensor); + // Plans ignore runtime bindings; allocator reservations must include them. + std::vector graph_layout(ggml_cgraph* graph, bool include_bindings); bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan); ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index); ggml_tensor* input_tensor(ggml_cgraph* gf, const Segment::InputRef& input_ref); std::vector param_tensors(ggml_cgraph* gf, const Segment& segment); - std::unordered_set collect_future_input_names(ggml_cgraph* gf, - const Plan& plan, - size_t current_segment_index); ggml_cgraph* build_segment_graph(ggml_cgraph* gf, const Segment& segment, ggml_context** graph_ctx_out); @@ -109,23 +92,12 @@ namespace sd::ggml_graph_cut { ggml_cgraph* gf, const std::unordered_set& params_tensor_set, const char* log_desc); - Plan apply_max_vram_budget(ggml_cgraph* gf, - const Plan& base_plan, - size_t max_graph_vram_bytes, - ggml_backend_t backend, - const std::unordered_set& params_tensor_set, - const char* log_desc); Plan resolve_plan(ggml_backend_t backend, ggml_cgraph* gf, PlanCache* cache, - size_t max_graph_vram_bytes, const std::unordered_set& params_tensor_set, const char* log_desc); - // Mark leading segments resident after reserving streamed execution headroom. - void annotate_residency(Plan& plan, - size_t max_graph_vram_bytes, - bool prefetch_enabled); } // namespace sd::ggml_graph_cut #endif // __SD_CORE_GGML_GRAPH_CUT_H__ diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp new file mode 100644 index 00000000..c3d69fd7 --- /dev/null +++ b/src/core/ggml_runner.cpp @@ -0,0 +1,296 @@ +#include +#include +#include + +#include "core/ggml_extend.hpp" +#include "core/segment_graph_bindings.h" +#include "core/segment_weight_pipeline.h" + +using namespace sd; + +static size_t add_bytes(size_t a, size_t b) { + return b > SIZE_MAX - a ? SIZE_MAX : a + b; +} + +ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph, size_t direct_bytes) { + auto external_backend = [&](const ggml_tensor* tensor) -> ggml_backend_t { + if (!params_tensor_set_.count(tensor)) { + return nullptr; + } + auto placement = graph_cut_layer_split_assignments_.find(tensor); + return placement == graph_cut_layer_split_assignments_.end() ? runtime_backend : placement->second; + }; + auto assign_nodes = [&](ggml_backend_sched_t scheduler, ggml_cgraph* copy) { + pin_multi_device_nodes(scheduler, copy, graph); + }; + return workspace_.measure(graph, direct_bytes, external_backend, assign_nodes); +} + +std::vector GGMLRunner::memory_requests( + const std::vector& sizes, + size_t pending_cache_bytes) const { + std::vector requests; + for (const auto& size : sizes) { + const size_t retained = retained_runtime_buffer_bytes(size.backend); + const size_t reusable = workspace_.bytes(size.backend); + const size_t cache_bytes = size.backend == runtime_backend ? pending_cache_bytes : 0; + const size_t pending = add_bytes(size.bytes > reusable ? size.bytes - reusable : 0, cache_bytes); + size_t limit = max_graph_vram_bytes; + if (is_multi_device()) { + size_t index = 0; + if (size.backend != runtime_backend) { + auto position = std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), size.backend); + index = static_cast(position - extra_runtime_backends.begin()) + 1; + } + if (index < graph_cut_layer_split_backend_vram_limits_.size()) { + limit = graph_cut_layer_split_backend_vram_limits_[index]; + } + } + requests.push_back({size.backend, reinterpret_cast(this), pending, + retained, limit}); + } + return requests; +} + +bool GGMLRunner::fits(const std::vector& requests, + const std::vector& params) const { + auto manager = residency_manager.lock(); + if (manager == nullptr) { + return params.empty(); + } + for (const auto& request : requests) { + if (!manager->fits_compute_backend_capacity(request, params)) { + return false; + } + } + return true; +} + +bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads) { + if (sd_backend_is_cpu(runtime_backend)) { + sd_backend_cpu_set_n_threads(runtime_backend, n_threads); + } + if (workspace_.cpu_backend() != nullptr) { + sd_backend_cpu_set_n_threads(workspace_.cpu_backend(), n_threads); + } + auto scheduler = workspace_.scheduler(); + ggml_status status; + if (scheduler != nullptr) { + if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) { + LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring", get_desc().c_str()); + multi_device_eval_callback_warned = true; + } + status = ggml_backend_sched_graph_compute(scheduler, graph); + } else { + status = sd_backend_graph_compute_with_eval_callback(runtime_backend, graph, + sd_get_backend_eval_callback(), + sd_get_backend_eval_callback_data()); + } + workspace_.synchronize(); + if (status != GGML_STATUS_SUCCESS) { + LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status)); + return false; + } + const std::string description = get_desc(); + if (!debug_tensors.empty()) { + std::unordered_set graph_tensors; + const int leaf_count = ggml_graph_cut::leaf_count(graph); + const int node_count = ggml_graph_n_nodes(graph); + graph_tensors.reserve(static_cast(leaf_count + node_count)); + for (int index = 0; index < leaf_count; ++index) { + graph_tensors.insert(ggml_graph_cut::leaf_tensor(graph, index)); + } + for (int index = 0; index < node_count; ++index) { + graph_tensors.insert(ggml_graph_node(graph, index)); + } + + for (const auto& entry : debug_tensors) { + ggml_tensor* tensor = entry.first; + if (tensor == nullptr || graph_tensors.find(tensor) == graph_tensors.end()) { + continue; + } + ggml_backend_buffer_t buffer = + tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer; + if (buffer == nullptr) { + LOG_WARN("%s skip debug tensor '%s': tensor buffer not set", + description.c_str(), + entry.second.c_str()); + continue; + } + if (tensor->type != GGML_TYPE_F32) { + LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s", + description.c_str(), + entry.second.c_str(), + ggml_type_name(tensor->type)); + continue; + } + auto debug_tensor = make_sd_tensor_from_ggml(tensor); + print_sd_tensor(debug_tensor, false, entry.second.c_str()); + } + } + + return true; +} + +std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function& read_outputs) { + if (!assign_graph_cut_layer_split_backends(graph)) { + return std::nullopt; + } + const auto params = collect_used_param_tensors(graph); + ggml_graph_cut::Plan plan; + if (!resolve_graph_cut_plan(graph, &plan)) { + return std::nullopt; + } + const auto full_measurement = measure(graph, plan.compute_buffer_size); + if (full_measurement.buffers.empty()) { + return std::nullopt; + } + auto manager = residency_manager.lock(); + const bool segmented = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) && + manager != nullptr && manager->segmented_compute_enabled() && + plan.valid && plan.has_cuts && plan.segments.size() > 1 && + !fits(memory_requests(full_measurement.buffers, cache_.pending_bytes(graph)), params); + if (!segmented) { + ggml_graph_cut::Segment segment; + segment.group_name = "graph"; + segment.compute_buffer_size = plan.compute_buffer_size; + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + segment.internal_node_indices.push_back(i); + } + for (int i = 0; i < ggml_graph_cut::leaf_count(graph); ++i) { + auto tensor = ggml_graph_cut::leaf_tensor(graph, i); + ggml_graph_cut::Segment::InputRef input; + input.leaf_index = i; + input.type = canonical_param_tensor(tensor) != nullptr + ? ggml_graph_cut::Segment::INPUT_PARAM + : ggml_graph_cut::Segment::INPUT_EXTERNAL; + segment.input_refs.push_back(input); + } + plan.segments = {std::move(segment)}; + } + const bool segments_changed = plan.segments.size() != logged_segment_count_; + if (segments_changed && (segmented || logged_segment_count_ > 1)) { + LOG_DEBUG("%s using %zu segment%s", get_desc().c_str(), + plan.segments.size(), plan.segments.size() == 1 ? "" : "s"); + } + SegmentGraphBindings bindings(cut_cache_, plan, graph); + SegmentWeightPipeline weights(manager, runtime_backend, reinterpret_cast(this), + graph, plan, params_tensor_set_, + segmented && manager != nullptr && manager->prefetch_enabled()); + + std::map peak_compute_bytes; + auto track_compute_buffer = [&](ggml_backend_t backend) { + if (backend != nullptr) { + auto& peak = peak_compute_bytes[backend]; + peak = std::max(peak, workspace_.bytes(backend)); + } + }; + std::optional> output = Tensor(); + for (size_t index = 0; index < plan.segments.size(); ++index) { + const auto& segment = plan.segments[index]; + const bool last = index + 1 == plan.segments.size(); + auto fail_segment = [&](const char* phase) { + LOG_ERROR("%s segment %zu/%zu (%s) failed during %s", get_desc().c_str(), + index + 1, plan.segments.size(), segment.group_name.c_str(), phase); + return std::nullopt; + }; + cut_cache_.prune(segment.live_cut_names); + bindings.reset(segment); + if (!bindings.bind_cached_inputs(segment, get_desc().c_str())) { + return fail_segment("input binding"); + } + ggml_context* segment_context = nullptr; + auto segment_graph = segmented + ? ggml_graph_cut::build_segment_graph(graph, segment, &segment_context) + : graph; + struct SegmentCleanup { + GGMLRunner& runner; + SegmentWeightPipeline& weights; + SegmentGraphBindings& bindings; + ggml_context* context; + ~SegmentCleanup() { + runner.workspace_.segment_end(); + bindings.restore(); + weights.segment_end(); + ggml_free(context); + runner.sync_runtime_residency(); + } + } segment_cleanup{*this, weights, bindings, segment_context}; + + auto measurement = segmented ? measure(segment_graph, segment.compute_buffer_size) : full_measurement; + if (!workspace_.prepare(measurement)) { + return fail_segment("workspace preparation"); + } + const size_t cut_bytes = last ? 0 : cut_cache_.estimate_output_bytes(graph, segment); + const size_t new_cache_bytes = add_bytes(cut_bytes, cache_.pending_bytes(segment_graph)); + auto ensure_capacity = [&]() { + sync_runtime_residency(); + auto requests = memory_requests(measurement.buffers, new_cache_bytes); + if (!fits(requests, weights.params(index)) && workspace_.release_excess(measurement)) { + sync_runtime_residency(); + requests = memory_requests(measurement.buffers, new_cache_bytes); + } + return weights.ensure_segment_capacity(index, requests); + }; + if (!weights.segment_start(index, ensure_capacity)) { + return fail_segment("weight preparation"); + } + // Preparing weights can execute LoRA graphs and reclaim an idle workspace. + if (!workspace_.measurement_matches(segment_graph, measurement)) { + measurement = measure(segment_graph, segment.compute_buffer_size); + } + if (!workspace_.prepare(measurement) || !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); + })) { + return fail_segment("workspace allocation"); + } + for (const auto& size : measurement.buffers) { + track_compute_buffer(size.backend); + } + if (workspace_.scheduler() != nullptr) { + track_compute_buffer(workspace_.cpu_backend()); + } + if (!ensure_capacity()) { + return fail_segment("allocated capacity check"); + } + copy_data_to_backend_tensor(segment_graph, false); + auto prefetch_requests = memory_requests(measurement.buffers, new_cache_bytes); + if (!prefetch_requests.empty()) { + weights.enqueue_next(index, prefetch_requests.front()); + } + if (!execute_segment(segment_graph, n_threads) || + !cache_.capture(segment_graph) || + !cut_cache_.capture(graph, segment, get_desc().c_str())) { + return fail_segment("execution or output caching"); + } + sync_runtime_residency(); + if (last) { + if (read_outputs && !read_outputs()) { + return fail_segment("output finalization"); + } + if (!no_return) { + auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str()); + output = read_graph_tensor(result, "output"); + if (!output.has_value()) { + return fail_segment("output readback"); + } + } + } + // Final outputs and their callbacks may still be views of consumed cuts. + cut_cache_.prune(segment.future_cut_names); + } + if (segments_changed || peak_compute_bytes != logged_compute_bytes_) { + for (const auto& entry : peak_compute_bytes) { + LOG_DEBUG("%s compute buffer size: %.2f MB(%s) on %s (peak across %zu segment%s)", + get_desc().c_str(), entry.second / (1024.0 * 1024.0), + sd_backend_is_cpu(entry.first) ? "RAM" : "VRAM", ggml_backend_name(entry.first), + plan.segments.size(), plan.segments.size() == 1 ? "" : "s"); + } + logged_compute_bytes_ = std::move(peak_compute_bytes); + logged_segment_count_ = plan.segments.size(); + } + return output; +} diff --git a/src/core/layer_split_partition.cpp b/src/core/layer_split_partition.cpp index 654b3056..f4059b04 100644 --- a/src/core/layer_split_partition.cpp +++ b/src/core/layer_split_partition.cpp @@ -145,19 +145,24 @@ namespace sd { std::vector backend_capacities = graph_cut_layer_split_backend_capacities(split_backends, backend_vram_limits, primary_backend_vram_limit); + // Existing placements may already occupy the reported free VRAM. Reuse + // them; execution checks missing weights and reclaims memory as needed. + const bool reuse_assignments = std::all_of(seen_params.begin(), seen_params.end(), [&](ggml_tensor* param) { + return param_assignments.count(param) != 0; + }); std::vector backend_by_segment(plan.segments.size(), split_backends[0]); size_t current_backend = 0; int64_t current_used = 0; for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) { int64_t bytes = segment_param_bytes[seg_idx]; - while (current_backend + 1 < split_backends.size() && + while (!reuse_assignments && current_backend + 1 < split_backends.size() && bytes > 0 && current_used + bytes > backend_capacities[current_backend]) { current_backend++; current_used = 0; } - if (bytes > 0 && current_used + bytes > backend_capacities[current_backend]) { + if (!reuse_assignments && bytes > 0 && current_used + bytes > backend_capacities[current_backend]) { LOG_ERROR("%s graph-cut layer split: segment %zu needs %.1f MB on %s, but only %.1f MB is available under current VRAM limits", desc, seg_idx, @@ -167,7 +172,6 @@ namespace sd { return false; } current_used += bytes; - backend_by_segment[seg_idx] = split_backends[current_backend]; for (ggml_tensor* param : segment_params[seg_idx]) { ggml_backend_t target_backend = split_backends[current_backend]; @@ -186,12 +190,16 @@ namespace sd { ggml_get_name(param)); return false; } - size_t backend_idx = (size_t)std::distance(split_backends.begin(), backend_it); + size_t backend_idx = (size_t)std::distance(split_backends.begin(), backend_it); + if (reuse_assignments) { + current_backend = backend_idx; + } assignment.first_segment_by_backend[backend_idx] = std::min(assignment.first_segment_by_backend[backend_idx], seg_idx); assignment.last_segment_by_backend[backend_idx] = std::max(assignment.last_segment_by_backend[backend_idx], seg_idx + 1); assignment.tensors_by_backend[backend_idx].push_back(param); assignment.bytes_by_backend[backend_idx] += (int64_t)ggml_nbytes(param); } + backend_by_segment[seg_idx] = split_backends[current_backend]; } const int n_nodes = ggml_graph_n_nodes(gf); diff --git a/src/core/layer_stream_prefetch.cpp b/src/core/layer_stream_prefetch.cpp deleted file mode 100644 index 5b737be7..00000000 --- a/src/core/layer_stream_prefetch.cpp +++ /dev/null @@ -1,124 +0,0 @@ -#include "core/layer_stream_prefetch.h" - -#include - -#include "core/ggml_graph_cut.h" -#include "weight_manager.h" - -namespace sd { - static ggml_tensor* canonical_param( - ggml_tensor* tensor, - const std::unordered_set& params) { - for (ggml_tensor* current = tensor; current != nullptr; current = current->view_src) { - if (params.find(current) != params.end()) { - return current; - } - } - return nullptr; - } - - LayerStreamPrefetch::LayerStreamPrefetch( - const std::shared_ptr& manager, - uintptr_t owner_id, - ggml_cgraph* graph, - const ggml_graph_cut::Plan& plan, - const std::unordered_set& params, - bool enabled) - : manager_(manager), - owner_id_(owner_id), - enabled_(enabled && manager != nullptr) { - segment_params_.resize(plan.segments.size()); - for (size_t segment_index = 0; segment_index < plan.segments.size(); ++segment_index) { - std::unordered_set seen; - for (ggml_tensor* tensor : - ggml_graph_cut::param_tensors(graph, plan.segments[segment_index])) { - ggml_tensor* param = canonical_param(tensor, params); - if (param != nullptr && seen.insert(param).second) { - segment_params_[segment_index].push_back(param); - } - } - } - } - - LayerStreamPrefetch::~LayerStreamPrefetch() { - clear(); - } - - size_t LayerStreamPrefetch::next_parameter_segment(size_t segment_index) const { - for (size_t next = segment_index + 1; next < segment_params_.size(); ++next) { - if (!segment_params_[next].empty()) { - return next; - } - } - return SIZE_MAX; - } - - void LayerStreamPrefetch::disable() { - clear(); - enabled_ = false; - } - - bool LayerStreamPrefetch::activate(size_t segment_index) { - if (!enabled_ || queued_segment_ == SIZE_MAX) { - return true; - } - if (queued_segment_ != segment_index) { - return true; - } - - auto manager = manager_.lock(); - if (manager == nullptr || - !manager->activate_prefetched_params(owner_id_, queued_params_)) { - disable(); - return false; - } - queued_params_.clear(); - queued_segment_ = SIZE_MAX; - return true; - } - - bool LayerStreamPrefetch::enqueue_next(size_t segment_index) { - if (!enabled_) { - return true; - } - if (queued_segment_ != SIZE_MAX) { - return true; - } - - const size_t next_segment = next_parameter_segment(segment_index); - if (next_segment == SIZE_MAX) { - return true; - } - - std::unordered_set active_params( - segment_params_[segment_index].begin(), - segment_params_[segment_index].end()); - std::vector params; - params.reserve(segment_params_[next_segment].size()); - for (ggml_tensor* param : segment_params_[next_segment]) { - if (active_params.find(param) == active_params.end()) { - params.push_back(param); - } - } - if (params.empty()) { - return true; - } - - auto manager = manager_.lock(); - if (manager == nullptr || !manager->prefetch_params(owner_id_, params)) { - disable(); - return false; - } - queued_params_ = std::move(params); - queued_segment_ = next_segment; - return true; - } - - void LayerStreamPrefetch::clear() { - if (auto manager = manager_.lock()) { - manager->clear_prefetched_params(owner_id_); - } - queued_params_.clear(); - queued_segment_ = SIZE_MAX; - } -} diff --git a/src/core/layer_stream_prefetch.h b/src/core/layer_stream_prefetch.h deleted file mode 100644 index 2457a1c6..00000000 --- a/src/core/layer_stream_prefetch.h +++ /dev/null @@ -1,48 +0,0 @@ -#ifndef __SD_CORE_LAYER_STREAM_PREFETCH_H__ -#define __SD_CORE_LAYER_STREAM_PREFETCH_H__ - -#include -#include -#include -#include -#include - -struct ggml_cgraph; -struct ggml_tensor; -struct RunnerWeightManager; - -namespace sd::ggml_graph_cut { - struct Plan; -} - -namespace sd { - class LayerStreamPrefetch { - private: - std::weak_ptr manager_; - uintptr_t owner_id_ = 0; - std::vector> segment_params_; - std::vector queued_params_; - size_t queued_segment_ = SIZE_MAX; - bool enabled_ = true; - - size_t next_parameter_segment(size_t segment_index) const; - void disable(); - - public: - LayerStreamPrefetch( - const std::shared_ptr& manager, - uintptr_t owner_id, - ggml_cgraph* graph, - const ggml_graph_cut::Plan& plan, - const std::unordered_set& params, - bool enabled = true); - ~LayerStreamPrefetch(); - - bool enabled() const { return enabled_; } - bool activate(size_t segment_index); - bool enqueue_next(size_t segment_index); - void clear(); - }; -} - -#endif // __SD_CORE_LAYER_STREAM_PREFETCH_H__ diff --git a/src/core/runner_cache.cpp b/src/core/runner_cache.cpp new file mode 100644 index 00000000..09562106 --- /dev/null +++ b/src/core/runner_cache.cpp @@ -0,0 +1,211 @@ +#include "core/runner_cache.h" + +#include +#include +#include + +#include "core/ggml_graph_cut.h" +#include "core/util.h" + +namespace sd { + static std::unordered_set cache_graph_tensors(ggml_cgraph* graph) { + std::unordered_set tensors; + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + tensors.insert(ggml_graph_node(graph, i)); + } + for (int i = 0; i < ggml_graph_cut::leaf_count(graph); ++i) { + tensors.insert(ggml_graph_cut::leaf_tensor(graph, i)); + } + return tensors; + } + + CachedTensor::~CachedTensor() { + ggml_backend_buffer_free(buffer); + ggml_free(context); + } + + std::unique_ptr CachedTensor::copy(ggml_backend_t backend, + const std::string& name, + ggml_tensor* source) { + if (ggml_graph_cut::tensor_buffer(source) == nullptr) { + return nullptr; + } + auto entry = std::make_unique(); + entry->context = ggml_init({2 * ggml_tensor_overhead(), nullptr, true}); + if (entry->context == nullptr) { + return nullptr; + } + entry->tensor = ggml_dup_tensor(entry->context, source); + // Cut views are rebound with their original strides and offsets. + std::copy(std::begin(source->nb), std::end(source->nb), std::begin(entry->tensor->nb)); + ggml_set_name(entry->tensor, name.c_str()); + entry->buffer = ggml_backend_alloc_ctx_tensors(entry->context, backend); + if (entry->buffer == nullptr) { + return nullptr; + } + if (source->view_src != nullptr || !ggml_is_contiguous(source) || source->buffer == nullptr) { + std::vector data(ggml_nbytes(source)); + ggml_backend_tensor_get(source, data.data(), 0, data.size()); + ggml_backend_tensor_set(entry->tensor, data.data(), 0, data.size()); + } else { + ggml_backend_tensor_copy(source, entry->tensor); + } + return entry; + } + + static ggml_tensor* cached_tensor(const CachedTensors& tensors, const std::string& name) { + auto entry = tensors.find(name); + return entry == tensors.end() ? nullptr : entry->second->tensor; + } + + static size_t resident_bytes(const CachedTensors& tensors, ggml_backend_dev_t device) { + size_t bytes = 0; + for (const auto& entry : tensors) { + auto buffer = entry.second->buffer; + if (!ggml_backend_buffer_is_host(buffer) && + ggml_backend_buft_get_device(ggml_backend_buffer_get_type(buffer)) == device) { + const size_t size = ggml_backend_buffer_get_size(buffer); + bytes = size > SIZE_MAX - bytes ? SIZE_MAX : bytes + size; + } + } + return bytes; + } + + ggml_tensor* RunnerCache::get(const std::string& name) const { + return cached_tensor(committed_, name); + } + + void RunnerCache::stage(const std::string& name, ggml_tensor* tensor) { + if (tensor != nullptr) { + ggml_set_output(tensor); + outputs_[name] = tensor; + } + } + + size_t RunnerCache::pending_bytes(ggml_cgraph* graph) const { + if (outputs_.empty()) { + return 0; + } + auto tensors = cache_graph_tensors(graph); + auto buft = ggml_backend_get_default_buffer_type(backend_); + size_t bytes = 0; + for (const auto& output : outputs_) { + if (pending_.count(output.first) || !tensors.count(output.second)) { + continue; + } + const size_t size = GGML_PAD(ggml_backend_buft_get_alloc_size(buft, output.second), + ggml_backend_buft_get_alignment(buft)); + bytes = size > SIZE_MAX - bytes ? SIZE_MAX : bytes + size; + } + return bytes; + } + + size_t RunnerCache::resident_bytes(ggml_backend_dev_t device) const { + const size_t committed = sd::resident_bytes(committed_, device); + const size_t pending = sd::resident_bytes(pending_, device); + return pending > SIZE_MAX - committed ? SIZE_MAX : committed + pending; + } + + bool RunnerCache::capture(ggml_cgraph* graph) { + if (outputs_.empty()) { + return true; + } + const auto tensors = cache_graph_tensors(graph); + for (const auto& output : outputs_) { + if (pending_.count(output.first) || !tensors.count(output.second)) { + continue; + } + GGML_ASSERT(ggml_is_contiguous(output.second)); + auto entry = CachedTensor::copy(backend_, output.first, output.second); + if (entry == nullptr) { + return false; + } + pending_[output.first] = std::move(entry); + } + ggml_backend_synchronize(backend_); + return true; + } + + void RunnerCache::graph_end(bool success) { + // Graph inputs can still reference the previous generation until graph end. + if (success) { + for (auto& entry : pending_) { + committed_[entry.first] = std::move(entry.second); + } + } + pending_.clear(); + outputs_.clear(); + } + + void RunnerCache::clear() { + graph_end(false); + committed_.clear(); + } + + ggml_tensor* GraphCutTensorCache::get(const std::string& name) const { + return cached_tensor(tensors_, name); + } + + size_t GraphCutTensorCache::resident_bytes(ggml_backend_dev_t device) const { + return sd::resident_bytes(tensors_, device); + } + + size_t GraphCutTensorCache::estimate_output_bytes( + ggml_cgraph* graph, + const ggml_graph_cut::Segment& segment) const { + ggml_backend_buffer_type_t buffer_type = + ggml_backend_get_default_buffer_type(backend_); + if (buffer_type == nullptr) { + return SIZE_MAX; + } + const size_t alignment = ggml_backend_buft_get_alignment(buffer_type); + size_t total_size = 0; + for (size_t output_idx = 0; output_idx < segment.output_node_indices.size(); ++output_idx) { + ggml_tensor* output = ggml_graph_cut::output_tensor(graph, segment, output_idx); + if (output == nullptr || !ggml_graph_cut::is_graph_cut_tensor(output) || + !segment.future_cut_names.count(output->name)) { + continue; + } + ggml_tensor* source = ggml_graph_cut::cache_source_tensor(output); + const size_t tensor_size = GGML_PAD( + ggml_backend_buft_get_alloc_size(buffer_type, source), alignment); + total_size = tensor_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + tensor_size; + } + return total_size; + } + + void GraphCutTensorCache::prune(const std::unordered_set& keep_names) { + for (auto it = tensors_.begin(); it != tensors_.end();) { + it = keep_names.count(it->first) ? std::next(it) : tensors_.erase(it); + } + } + + bool GraphCutTensorCache::capture(ggml_cgraph* graph, + const ggml_graph_cut::Segment& segment, + const char* log_desc) { + size_t copied_bytes = 0; + size_t copied_count = 0; + for (int index : segment.output_node_indices) { + auto output = ggml_graph_node(graph, index); + if (!ggml_graph_cut::is_graph_cut_tensor(output) || + !segment.future_cut_names.count(output->name)) { + continue; + } + auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output)); + if (entry == nullptr) { + LOG_ERROR("%s failed to capture graph cut tensor: %s", log_desc, output->name); + return false; + } + const size_t size = ggml_backend_buffer_get_size(entry->buffer); + copied_bytes = size > SIZE_MAX - copied_bytes ? SIZE_MAX : copied_bytes + size; + ++copied_count; + tensors_[output->name] = std::move(entry); + } + ggml_backend_synchronize(backend_); + if (copied_count > 0) { + LOG_DEBUG("%s graph cut cache added %6.2f MB (%zu tensors)", + log_desc, copied_bytes / (1024.f * 1024.f), copied_count); + } + return true; + } +} diff --git a/src/core/runner_cache.h b/src/core/runner_cache.h new file mode 100644 index 00000000..bd331a4b --- /dev/null +++ b/src/core/runner_cache.h @@ -0,0 +1,66 @@ +#ifndef __SD_CORE_RUNNER_CACHE_H__ +#define __SD_CORE_RUNNER_CACHE_H__ + +#include +#include +#include +#include + +#include "ggml-backend.h" + +namespace sd::ggml_graph_cut { + struct Segment; +} + +namespace sd { + struct CachedTensor { + ggml_context* context = nullptr; + ggml_backend_buffer_t buffer = nullptr; + ggml_tensor* tensor = nullptr; + ~CachedTensor(); + static std::unique_ptr copy(ggml_backend_t backend, + const std::string& name, + ggml_tensor* source); + }; + using CachedTensors = std::map>; + + class RunnerCache { + ggml_backend_t backend_; + CachedTensors committed_; + CachedTensors pending_; + std::map outputs_; + + public: + explicit RunnerCache(ggml_backend_t backend) + : backend_(backend) {} + RunnerCache(const RunnerCache&) = delete; + RunnerCache& operator=(const RunnerCache&) = delete; + + ggml_tensor* get(const std::string& name) const; + void stage(const std::string& name, ggml_tensor* tensor); + const std::map& outputs() const { return outputs_; } + size_t pending_bytes(ggml_cgraph* graph) const; + size_t resident_bytes(ggml_backend_dev_t device) const; + bool capture(ggml_cgraph* graph); + void graph_end(bool success); + void clear(); + }; + + class GraphCutTensorCache { + ggml_backend_t backend_; + CachedTensors tensors_; + + public: + explicit GraphCutTensorCache(ggml_backend_t backend) + : backend_(backend) {} + ggml_tensor* get(const std::string& name) const; + size_t resident_bytes(ggml_backend_dev_t device) const; + size_t estimate_output_bytes(ggml_cgraph* graph, + const ggml_graph_cut::Segment& segment) const; + bool capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc); + void prune(const std::unordered_set& keep_names); + void clear() { tensors_.clear(); } + }; +} + +#endif // __SD_CORE_RUNNER_CACHE_H__ diff --git a/src/core/segment_graph_bindings.cpp b/src/core/segment_graph_bindings.cpp new file mode 100644 index 00000000..3c40a3b0 --- /dev/null +++ b/src/core/segment_graph_bindings.cpp @@ -0,0 +1,137 @@ +#include "core/segment_graph_bindings.h" + +#include +#include + +#include "core/ggml_graph_cut.h" +#include "core/runner_cache.h" +#include "core/util.h" +#include "ggml.h" + +namespace sd { + SegmentGraphBindings::SegmentGraphBindings(GraphCutTensorCache& tensor_cache, + const ggml_graph_cut::Plan& plan, + ggml_cgraph* graph) + : tensor_cache_(tensor_cache), + graph_(graph) { + GGML_ASSERT(graph_ != nullptr); + for (int i = 0; i < ggml_graph_n_nodes(graph_); ++i) { + ggml_tensor* tensor = ggml_graph_node(graph_, i); + Topology topology{tensor->op, {}, tensor->view_src, tensor->flags}; + std::copy(std::begin(tensor->src), std::end(tensor->src), topology.sources.begin()); + topology_[tensor] = topology; + } + for (const auto& segment : plan.segments) { + for (const auto& input : segment.input_refs) { + if (input.type != ggml_graph_cut::Segment::INPUT_EXTERNAL) { + continue; + } + ggml_tensor* tensor = ggml_graph_cut::input_tensor(graph_, input); + if (tensor == nullptr || tensor->buffer == nullptr) { + continue; + } + external_bindings_[tensor] = {tensor->buffer, tensor->data, tensor->extra}; + } + } + } + + void SegmentGraphBindings::reset(const ggml_graph_cut::Segment& segment) { + restore(); + for (const auto& input : segment.input_refs) { + ggml_tensor* tensor = ggml_graph_cut::input_tensor(graph_, input); + if (tensor == nullptr) { + continue; + } + switch (input.type) { + case ggml_graph_cut::Segment::INPUT_PREVIOUS_CUT: + tensor->buffer = nullptr; + tensor->data = nullptr; + tensor->extra = nullptr; + break; + case ggml_graph_cut::Segment::INPUT_EXTERNAL: { + auto binding = external_bindings_.find(tensor); + if (binding != external_bindings_.end()) { + tensor->buffer = binding->second.buffer; + tensor->data = binding->second.data; + tensor->extra = binding->second.extra; + } else { + tensor->buffer = nullptr; + tensor->data = nullptr; + tensor->extra = nullptr; + } + break; + } + case ggml_graph_cut::Segment::INPUT_PARAM: + break; + } + } + + for (int node_index : segment.internal_node_indices) { + ggml_tensor* node = ggml_graph_node(graph_, node_index); + if (node == nullptr) { + continue; + } + node->buffer = nullptr; + node->data = nullptr; + node->extra = nullptr; + } + } + + void SegmentGraphBindings::restore() { + for (const auto& entry : topology_) { + entry.first->op = entry.second.op; + entry.first->view_src = entry.second.view_source; + entry.first->flags = entry.second.flags; + std::copy(entry.second.sources.begin(), entry.second.sources.end(), std::begin(entry.first->src)); + } + } + + bool SegmentGraphBindings::bind_cached_inputs( + const ggml_graph_cut::Segment& segment, + const char* log_desc) { + std::unordered_map cached_view_sources; + for (const auto& input : segment.input_refs) { + if (input.type != ggml_graph_cut::Segment::INPUT_PREVIOUS_CUT) { + continue; + } + ggml_tensor* input_tensor = ggml_graph_cut::input_tensor(graph_, input); + if (input_tensor == nullptr) { + continue; + } + ggml_tensor* cached_tensor = tensor_cache_.get(input.display_name); + if (cached_tensor == nullptr) { + LOG_ERROR("%s missing graph cut cache tensor: %s", + log_desc, + input.display_name.c_str()); + return false; + } + if (input_tensor->view_src != nullptr) { + cached_view_sources[topology_.at(input_tensor).view_source] = cached_tensor; + input_tensor->view_src = cached_tensor; + input_tensor->buffer = nullptr; + input_tensor->data = cached_tensor->data == nullptr + ? nullptr + : static_cast(static_cast(cached_tensor->data) + + input_tensor->view_offs); + input_tensor->extra = cached_tensor->extra; + } else { + input_tensor->buffer = cached_tensor->buffer; + input_tensor->data = cached_tensor->data; + input_tensor->extra = cached_tensor->extra; + } + for (int source_index = 0; source_index < GGML_MAX_SRC; ++source_index) { + input_tensor->src[source_index] = nullptr; + } + input_tensor->op = GGML_OP_NONE; + } + // ggml flattens view chains, so descendants also need the cached root. + for (int node_index : segment.internal_node_indices) { + ggml_tensor* node = ggml_graph_node(graph_, node_index); + auto cached_source = cached_view_sources.find(topology_.at(node).view_source); + if (cached_source != cached_view_sources.end()) { + node->view_src = cached_source->second; + } + } + return true; + } +} diff --git a/src/core/segment_graph_bindings.h b/src/core/segment_graph_bindings.h new file mode 100644 index 00000000..08aba156 --- /dev/null +++ b/src/core/segment_graph_bindings.h @@ -0,0 +1,52 @@ +#ifndef __SD_CORE_SEGMENT_GRAPH_BINDINGS_H__ +#define __SD_CORE_SEGMENT_GRAPH_BINDINGS_H__ + +#include +#include +#include "ggml.h" + +struct ggml_backend_buffer; +struct ggml_cgraph; +struct ggml_tensor; + +namespace sd { + class GraphCutTensorCache; + + namespace ggml_graph_cut { + struct Plan; + struct Segment; + } + + class SegmentGraphBindings { + public: + SegmentGraphBindings(GraphCutTensorCache& tensor_cache, + const ggml_graph_cut::Plan& plan, + ggml_cgraph* graph); + + void reset(const ggml_graph_cut::Segment& segment); + void restore(); + ~SegmentGraphBindings() { restore(); } + bool bind_cached_inputs(const ggml_graph_cut::Segment& segment, + const char* log_desc); + + private: + struct ExternalBinding { + ggml_backend_buffer* buffer = nullptr; + void* data = nullptr; + void* extra = nullptr; + }; + + GraphCutTensorCache& tensor_cache_; + ggml_cgraph* graph_ = nullptr; + std::unordered_map external_bindings_; + struct Topology { + ggml_op op; + std::array sources; + ggml_tensor* view_source; + int flags; + }; + std::unordered_map topology_; + }; +} + +#endif // __SD_CORE_SEGMENT_GRAPH_BINDINGS_H__ diff --git a/src/core/segment_weight_pipeline.cpp b/src/core/segment_weight_pipeline.cpp new file mode 100644 index 00000000..f769794d --- /dev/null +++ b/src/core/segment_weight_pipeline.cpp @@ -0,0 +1,205 @@ +#include "core/segment_weight_pipeline.h" + +#include + +#include "core/ggml_graph_cut.h" +#include "device_residency_manager.h" + +namespace sd { + static ggml_tensor* canonical_param( + ggml_tensor* tensor, + const std::unordered_set& params) { + for (ggml_tensor* current = tensor; current != nullptr; current = current->view_src) { + if (params.find(current) != params.end()) { + return current; + } + } + return nullptr; + } + + SegmentWeightPipeline::SegmentWeightPipeline( + const std::shared_ptr& residency_manager, + ggml_backend_t compute_backend, + uintptr_t owner_id, + ggml_cgraph* graph, + const ggml_graph_cut::Plan& plan, + const std::unordered_set& params, + bool enabled) + : residency_manager_(residency_manager), + compute_backend_(compute_backend), + owner_id_(owner_id), + enabled_(enabled && residency_manager != nullptr) { + segment_params_.resize(plan.segments.size()); + for (size_t segment_index = 0; segment_index < plan.segments.size(); ++segment_index) { + std::unordered_set seen; + for (ggml_tensor* tensor : + ggml_graph_cut::param_tensors(graph, plan.segments[segment_index])) { + ggml_tensor* param = canonical_param(tensor, params); + if (param != nullptr && seen.insert(param).second) { + segment_params_[segment_index].push_back(param); + } + } + } + } + + SegmentWeightPipeline::~SegmentWeightPipeline() { + segment_end(); + clear(); + } + + size_t SegmentWeightPipeline::next_parameter_segment(size_t segment_index) const { + for (size_t next = segment_index + 1; next < segment_params_.size(); ++next) { + if (!segment_params_[next].empty()) { + return next; + } + } + return SIZE_MAX; + } + + std::vector> SegmentWeightPipeline::preferred_eviction_order() const { + return {segment_params_.rbegin(), segment_params_.rend()}; + } + + void SegmentWeightPipeline::disable() { + clear(); + enabled_ = false; + } + + void SegmentWeightPipeline::activate(size_t segment_index) { + if (!enabled_ || queued_segment_ == SIZE_MAX || queued_segment_ != segment_index) { + return; + } + + auto manager = residency_manager_.lock(); + if (manager == nullptr || + !manager->activate_prefetched_params(owner_id_, queued_params_)) { + disable(); + return; + } + queued_params_.clear(); + queued_segment_ = SIZE_MAX; + } + + bool SegmentWeightPipeline::ensure_segment_capacity( + size_t segment_index, + const std::vector& requests) { + if (segment_index >= segment_params_.size()) { + return false; + } + auto manager = residency_manager_.lock(); + if (manager == nullptr) { + return segment_params_[segment_index].empty(); + } + std::vector protected_params = segment_params_[segment_index]; + protected_params.insert(protected_params.end(), queued_params_.begin(), queued_params_.end()); + for (const auto& request : requests) { + if (!manager->ensure_compute_backend_capacity(request, segment_params_[segment_index], + preferred_eviction_order(), protected_params)) { + return false; + } + } + return true; + } + + bool SegmentWeightPipeline::segment_start(size_t segment_index, const std::function& ensure_capacity) { + GGML_ASSERT(pinned_params_.empty()); + activate(segment_index); + if (!ensure_capacity()) { + return false; + } + auto manager = residency_manager_.lock(); + if (manager == nullptr) { + return segment_params_[segment_index].empty(); + } + if (!manager->prepare_params(segment_params_[segment_index])) { + return false; + } + pinned_params_ = segment_params_[segment_index]; + return true; + } + + void SegmentWeightPipeline::segment_end() { + if (auto manager = residency_manager_.lock()) { + manager->release_compute_backend_params(pinned_params_); + } + pinned_params_.clear(); + } + + void SegmentWeightPipeline::enqueue_next( + size_t segment_index, + const DeviceMemoryRequest& request) { + if (!enabled_ || queued_segment_ != SIZE_MAX) { + return; + } + + const size_t next_segment = next_parameter_segment(segment_index); + if (next_segment == SIZE_MAX) { + return; + } + + std::unordered_set active_params( + segment_params_[segment_index].begin(), + segment_params_[segment_index].end()); + std::vector params; + params.reserve(segment_params_[next_segment].size()); + for (ggml_tensor* param : segment_params_[next_segment]) { + if (active_params.find(param) == active_params.end()) { + params.push_back(param); + } + } + if (params.empty()) { + return; + } + + auto manager = residency_manager_.lock(); + if (manager == nullptr) { + disable(); + return; + } + const WeightResidencyInfo residency = + manager->inspect_compute_backend_params(params); + if (residency.missing_bytes == 0) { + return; + } + if (!residency.async_prefetch_supported) { + disable(); + return; + } + + DeviceMemoryRequest backend_request = request; + backend_request.compute_backend = compute_backend_; + backend_request.owner_id = owner_id_; + std::vector protected_params = segment_params_[segment_index]; + protected_params.insert(protected_params.end(), params.begin(), params.end()); + if (!manager->ensure_compute_backend_capacity(backend_request, + params, + preferred_eviction_order(), + protected_params)) { + disable(); + return; + } + switch (manager->prefetch_params(owner_id_, params)) { + case WeightPrefetchResult::Scheduled: + queued_params_ = std::move(params); + queued_segment_ = next_segment; + return; + case WeightPrefetchResult::AlreadyResident: + return; + case WeightPrefetchResult::Unsupported: + disable(); + return; + case WeightPrefetchResult::Failed: + disable(); + return; + } + disable(); + } + + void SegmentWeightPipeline::clear() { + if (auto manager = residency_manager_.lock()) { + manager->clear_prefetched_params(owner_id_); + } + queued_params_.clear(); + queued_segment_ = SIZE_MAX; + } +} diff --git a/src/core/segment_weight_pipeline.h b/src/core/segment_weight_pipeline.h new file mode 100644 index 00000000..598bf9f4 --- /dev/null +++ b/src/core/segment_weight_pipeline.h @@ -0,0 +1,62 @@ +#ifndef __SD_CORE_SEGMENT_WEIGHT_PIPELINE_H__ +#define __SD_CORE_SEGMENT_WEIGHT_PIPELINE_H__ + +#include +#include +#include +#include +#include +#include + +#include "ggml-backend.h" + +struct DeviceMemoryRequest; +struct DeviceResidencyManager; +struct ggml_cgraph; +struct ggml_tensor; + +namespace sd::ggml_graph_cut { + struct Plan; +} + +namespace sd { + class SegmentWeightPipeline { + private: + std::weak_ptr residency_manager_; + ggml_backend_t compute_backend_ = nullptr; + uintptr_t owner_id_ = 0; + std::vector> segment_params_; + std::vector queued_params_; + std::vector pinned_params_; + size_t queued_segment_ = SIZE_MAX; + bool enabled_ = true; + + size_t next_parameter_segment(size_t segment_index) const; + std::vector> preferred_eviction_order() const; + void disable(); + void activate(size_t segment_index); + void clear(); + + public: + SegmentWeightPipeline( + const std::shared_ptr& residency_manager, + ggml_backend_t compute_backend, + uintptr_t owner_id, + ggml_cgraph* graph, + const ggml_graph_cut::Plan& plan, + const std::unordered_set& params, + bool enabled = true); + ~SegmentWeightPipeline(); + + const std::vector& params(size_t index) const { return segment_params_[index]; } + bool ensure_segment_capacity(size_t segment_index, + const std::vector& requests); + bool segment_start(size_t segment_index, const std::function& ensure_capacity); + void segment_end(); + // Prefetch is best effort; segment_start falls back to synchronous loading. + void enqueue_next(size_t segment_index, + const DeviceMemoryRequest& request); + }; +} + +#endif // __SD_CORE_SEGMENT_WEIGHT_PIPELINE_H__ diff --git a/src/core/tensor_ggml.hpp b/src/core/tensor_ggml.hpp index 774574f7..66ea2220 100644 --- a/src/core/tensor_ggml.hpp +++ b/src/core/tensor_ggml.hpp @@ -9,6 +9,7 @@ #include #include "core/tensor.hpp" +#include "ggml-backend.h" #include "ggml.h" namespace sd { @@ -54,10 +55,28 @@ namespace sd { GGML_ABORT("ggml tensor type does not match sd::Tensor type"); } Tensor result(shape_from_ggml(tensor)); - if (tensor->buffer != nullptr) { - ggml_backend_tensor_get(tensor, result.data(), 0, ggml_nbytes(tensor)); + std::vector strided_data; + void* destination = result.data(); + if (!ggml_is_contiguous(tensor)) { + strided_data.resize(ggml_nbytes(tensor)); + destination = strided_data.data(); + } + auto buffer = tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer; + if (buffer != nullptr) { + ggml_backend_tensor_get(tensor, destination, 0, ggml_nbytes(tensor)); } else { - std::memcpy(result.data(), tensor->data, ggml_nbytes(tensor)); + std::memcpy(destination, tensor->data, ggml_nbytes(tensor)); + } + if (!strided_data.empty()) { + for (int64_t i = 0; i < result.numel(); ++i) { + int64_t index = i; + size_t offset = 0; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + offset += static_cast(index % tensor->ne[d]) * tensor->nb[d]; + index /= tensor->ne[d]; + } + std::memcpy(result.data() + i, strided_data.data() + offset, sizeof(T)); + } } return result; } diff --git a/src/detailer.cpp b/src/detailer.cpp index fa9884be..2523dd51 100644 --- a/src/detailer.cpp +++ b/src/detailer.cpp @@ -715,7 +715,7 @@ std::vector ADetailerGGML::predict(sd_image_t image, LetterboxInput input = make_letterbox_input(image, params.input_size); int64_t start = ggml_time_ms(); sd::Tensor raw = detector->compute(n_threads, input.tensor); - detector->free_compute_buffer(); + detector->runner_end(); if (raw.empty()) { LOG_ERROR("YOLOv8 detector inference failed"); return {}; diff --git a/src/device_residency_manager.h b/src/device_residency_manager.h new file mode 100644 index 00000000..2fb2a342 --- /dev/null +++ b/src/device_residency_manager.h @@ -0,0 +1,76 @@ +#ifndef __DEVICE_RESIDENCY_MANAGER_H__ +#define __DEVICE_RESIDENCY_MANAGER_H__ + +#include +#include +#include + +#include "ggml-backend.h" + +struct ggml_tensor; + +enum class WeightPrefetchResult { + Scheduled, + AlreadyResident, + Unsupported, + Failed, +}; + +struct WeightResidencyInfo { + bool async_prefetch_supported = false; + size_t missing_bytes = 0; +}; + +struct DeviceMemoryRequest { + ggml_backend_t compute_backend = nullptr; + uintptr_t owner_id = 0; + size_t pending_allocation_bytes = 0; + size_t runtime_resident_bytes = 0; + size_t max_backend_bytes = 0; + + // Runtime buffers only; the manager accounts for weights separately. + size_t runtime_peak_bytes() const { + return pending_allocation_bytes > SIZE_MAX - runtime_resident_bytes + ? SIZE_MAX + : runtime_resident_bytes + pending_allocation_bytes; + } +}; + +struct DeviceResidencyManager { + virtual ~DeviceResidencyManager() = default; + + virtual bool segmented_compute_enabled() const = 0; + virtual bool prefetch_enabled() const = 0; + virtual void set_workspace_reclaimer(uintptr_t owner_id, std::function reclaim) = 0; + virtual void remove_runtime_owner(uintptr_t owner_id) = 0; + // Capacity requests select their backend's weights; protection spans all backends. + virtual bool fits_compute_backend_capacity(const DeviceMemoryRequest& request, + const std::vector& required_params) const = 0; + virtual bool assign_compute_backend(const std::vector& tensors, + ggml_backend_t compute_backend) = 0; + virtual bool prepare_params(const std::vector& tensors) = 0; + virtual void release_compute_backend_params(const std::vector& tensors) = 0; + virtual void evict_compute_backend_params(const std::vector& tensors) = 0; + virtual WeightResidencyInfo inspect_compute_backend_params( + const std::vector& tensors) const = 0; + virtual void update_runtime_residency(uintptr_t owner_id, + ggml_backend_t compute_backend, + size_t resident_bytes) = 0; + virtual bool ensure_compute_backend_capacity( + const DeviceMemoryRequest& request, + const std::vector& required_params, + const std::vector>& preferred_eviction_order, + const std::vector& protected_params) = 0; + virtual WeightPrefetchResult prefetch_params( + uintptr_t owner_id, + const std::vector& tensors) = 0; + virtual bool activate_prefetched_params(uintptr_t owner_id, + const std::vector& tensors) = 0; + virtual void clear_prefetched_params(uintptr_t owner_id) = 0; +}; + +// Transitional alias for model constructors that have not yet adopted the +// residency-oriented name. It does not introduce a second implementation. +using RunnerWeightManager = DeviceResidencyManager; + +#endif // __DEVICE_RESIDENCY_MANAGER_H__ diff --git a/src/extensions/generation_extension.h b/src/extensions/generation_extension.h index 67085c15..b1ed9729 100644 --- a/src/extensions/generation_extension.h +++ b/src/extensions/generation_extension.h @@ -49,7 +49,7 @@ struct GenerationExtension { virtual void get_param_tensors(std::map&) {} virtual void collect_loras(std::vector&) {} virtual void add_ignore_tensors(std::set&) const {} - virtual void runner_done() {} + virtual void runner_end() {} virtual void reset_runtime_condition() {} virtual bool prepare_condition(GenerationExtensionConditionContext&) { return false; diff --git a/src/extensions/photomaker_extension.cpp b/src/extensions/photomaker_extension.cpp index 48b023f9..ecf0081c 100644 --- a/src/extensions/photomaker_extension.cpp +++ b/src/extensions/photomaker_extension.cpp @@ -175,9 +175,9 @@ struct PhotoMakerExtension : public GenerationExtension { ignore_tensors.insert("pmid.unet."); } - void runner_done() override { + void runner_end() override { if (pmid_model != nullptr) { - pmid_model->runner_done(); + pmid_model->runner_end(); } } diff --git a/src/model/adapter/ip_adapter.hpp b/src/model/adapter/ip_adapter.hpp index cc4933e7..94f623ad 100644 --- a/src/model/adapter/ip_adapter.hpp +++ b/src/model/adapter/ip_adapter.hpp @@ -200,7 +200,7 @@ namespace IPAdapter { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(image_embeds); }; - return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true, true, true)); + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true)); } }; diff --git a/src/model/adapter/lora.hpp b/src/model/adapter/lora.hpp index 8eb88d54..4d701c90 100644 --- a/src/model/adapter/lora.hpp +++ b/src/model/adapter/lora.hpp @@ -125,13 +125,12 @@ struct LoraModel : public GGMLRunner { } void release_loaded_tensors() { - runner_done(); - free_compute_buffer(); + runner_end(); model_manager.reset(); free_params_ctx(); alloc_params_ctx(); - model_manager = std::make_shared(); - weight_manager = model_manager; + model_manager = std::make_shared(); + residency_manager = model_manager; lora_tensors.clear(); original_tensor_to_final_tensor.clear(); applied_lora_tensors.clear(); @@ -952,16 +951,19 @@ struct LoraModel : public GGMLRunner { auto get_graph = [&]() -> ggml_cgraph* { return build_lora_graph(model_tensors, model_tensor_names, version); }; - GGMLRunner::compute(get_graph, n_threads, false, false, false, true); - stat(!warn_unused); - for (auto item : original_tensor_to_final_tensor) { - ggml_tensor* original_tensor = item.first; - ggml_tensor* final_tensor = item.second; - - ggml_backend_tensor_copy(final_tensor, original_tensor); + auto read_outputs = [&]() { + for (const auto& item : original_tensor_to_final_tensor) { + ggml_backend_tensor_copy(item.second, item.first); + } + return true; + }; + auto result = GGMLRunner::compute(get_graph, n_threads, false, true, read_outputs); + if (!result.has_value()) { + LOG_ERROR("LoRA graph execution failed"); } + stat(!warn_unused); original_tensor_to_final_tensor.clear(); - GGMLRunner::free_compute_buffer(); + runner_end(); } void apply(std::map model_tensors, SDVersion version, int n_threads, bool warn_unused = true) { diff --git a/src/model/adapter/pmid.hpp b/src/model/adapter/pmid.hpp index 8f7d4dbd..58e90756 100644 --- a/src/model/adapter/pmid.hpp +++ b/src/model/adapter/pmid.hpp @@ -558,7 +558,7 @@ public: return build_graph(id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds); }; - return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true, true, true)); + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true)); } }; diff --git a/src/model/detector/yolov8.h b/src/model/detector/yolov8.h index a90fdf9a..3e9cd3b0 100644 --- a/src/model/detector/yolov8.h +++ b/src/model/detector/yolov8.h @@ -355,7 +355,7 @@ struct YOLOv8Runner : public GGMLRunner { sd::Tensor compute(int n_threads, const sd::Tensor& input) { auto get_graph = [&]() { return build_graph(input); }; - return take_or_empty(GGMLRunner::compute(get_graph, n_threads, false, false, false)); + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, false)); } }; diff --git a/src/model/diffusion/anima.hpp b/src/model/diffusion/anima.hpp index 4fe7e465..867f397b 100644 --- a/src/model/diffusion/anima.hpp +++ b/src/model/diffusion/anima.hpp @@ -717,7 +717,7 @@ namespace Anima { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, t5_ids, t5_weights, ref_latents); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/src/model/diffusion/boogu.hpp b/src/model/diffusion/boogu.hpp index 7bdbcd81..b199208d 100644 --- a/src/model/diffusion/boogu.hpp +++ b/src/model/diffusion/boogu.hpp @@ -815,7 +815,7 @@ namespace Boogu { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, ref_latents); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/src/model/diffusion/control.hpp b/src/model/diffusion/control.hpp index bf3c7e43..c8e72d7c 100644 --- a/src/model/diffusion/control.hpp +++ b/src/model/diffusion/control.hpp @@ -423,19 +423,26 @@ struct ControlNet : public GGMLRunner { return build_graph(x, hint, timesteps, context, y); }; - auto compute_result = GGMLRunner::compute(get_graph, n_threads, false, false, false, true); + auto read_outputs = [&]() { + controls.clear(); + controls.reserve(control_outputs_ggml.size()); + for (ggml_tensor* control : control_outputs_ggml) { + auto control_host = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml(control), 4); + if (control_host.empty()) { + return false; + } + controls.push_back(std::move(control_host)); + } + return true; + }; + auto compute_result = GGMLRunner::compute(get_graph, n_threads, false, true, read_outputs); + control_outputs_ggml.clear(); + guided_hint_output_ggml = nullptr; if (!compute_result.has_value()) { + controls.clear(); return std::nullopt; } - guided_hint_cached = get_cache_tensor_by_name(guided_hint_cache_name()) != nullptr; - controls.clear(); - controls.reserve(control_outputs_ggml.size()); - for (ggml_tensor* control : control_outputs_ggml) { - auto control_host = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml(control), 4); - GGML_ASSERT(!control_host.empty()); - controls.push_back(std::move(control_host)); - } return controls; } @@ -444,10 +451,10 @@ struct ControlNet : public GGMLRunner { std::map tensors; control_net.get_param_tensors(tensors); - auto manager = std::dynamic_pointer_cast(weight_manager.lock()); + auto manager = std::dynamic_pointer_cast(residency_manager.lock()); if (manager == nullptr) { owned_model_manager = std::make_shared(); - weight_manager = owned_model_manager; + residency_manager = owned_model_manager; manager = owned_model_manager; } diff --git a/src/model/diffusion/ernie_image.hpp b/src/model/diffusion/ernie_image.hpp index 12fcada5..4dfb7a89 100644 --- a/src/model/diffusion/ernie_image.hpp +++ b/src/model/diffusion/ernie_image.hpp @@ -440,7 +440,7 @@ namespace ErnieImage { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/src/model/diffusion/flux.hpp b/src/model/diffusion/flux.hpp index 375576c1..28a19a8a 100644 --- a/src/model/diffusion/flux.hpp +++ b/src/model/diffusion/flux.hpp @@ -1626,7 +1626,7 @@ namespace Flux { return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, ref_index_mode, skip_layers, pulid_id, pulid_id_weight); }; - auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); return result; } diff --git a/src/model/diffusion/hidream_o1.hpp b/src/model/diffusion/hidream_o1.hpp index 9d3df039..42d58d81 100644 --- a/src/model/diffusion/hidream_o1.hpp +++ b/src/model/diffusion/hidream_o1.hpp @@ -325,13 +325,11 @@ namespace HiDreamO1 { sd::Tensor compute(int n_threads, const sd::Tensor& image, - bool auto_free = true, - bool free_compute_buffer = true, - bool free_compute_params = true) { + bool auto_runner_end = true) { auto get_graph = [&]() { return build_graph(image); }; - auto output = GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params); + auto output = GGMLRunner::compute(get_graph, n_threads, auto_runner_end); return output.has_value() ? std::move(output.value()) : sd::Tensor(); } }; @@ -459,7 +457,7 @@ namespace HiDreamO1 { auto get_graph = [&]() { return build_graph(x, timestep, input_ids, input_pos, token_types, vinput_mask, image_embeds, ref_images); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, @@ -510,8 +508,8 @@ namespace HiDreamO1 { vision_runner->set_weight_adapter(adapter); } - void runner_done() override { - vision_runner->runner_done(); + void runner_end() override { + vision_runner->runner_end(); } SDCondition get_learned_condition(int n_threads, @@ -659,7 +657,7 @@ namespace HiDreamO1 { result.c_vinput_mask = sd::Tensor(vinput_mask_shape, std::move(vinput_mask)); result.c_image_embeds.reserve(vlm_images.size()); for (const auto& vlm_image : vlm_images) { - auto image_embed = vision_runner->compute(n_threads, vlm_image.second, false, true, true); + auto image_embed = vision_runner->compute(n_threads, vlm_image.second, false); if (image_embed.empty()) { LOG_ERROR("hidream_o1 conditioner: encode VLM image failed"); return SDCondition(); diff --git a/src/model/diffusion/hunyuan.hpp b/src/model/diffusion/hunyuan.hpp index f8108563..241c3560 100644 --- a/src/model/diffusion/hunyuan.hpp +++ b/src/model/diffusion/hunyuan.hpp @@ -654,7 +654,7 @@ namespace Hunyuan { return build_graph(x, timesteps, context, c_concat, y, guidance, byt5, vision, timestep_r); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/src/model/diffusion/ideogram4.hpp b/src/model/diffusion/ideogram4.hpp index 6ce8e1fa..1d2d9498 100644 --- a/src/model/diffusion/ideogram4.hpp +++ b/src/model/diffusion/ideogram4.hpp @@ -537,7 +537,7 @@ namespace Ideogram4 { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, use_uncond_model); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/src/model/diffusion/krea2.hpp b/src/model/diffusion/krea2.hpp index b3947b71..a80f5e42 100644 --- a/src/model/diffusion/krea2.hpp +++ b/src/model/diffusion/krea2.hpp @@ -775,7 +775,7 @@ namespace Krea2 { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, ref_latents, ref_image_params); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/src/model/diffusion/lens.hpp b/src/model/diffusion/lens.hpp index 931a8527..2658359c 100644 --- a/src/model/diffusion/lens.hpp +++ b/src/model/diffusion/lens.hpp @@ -408,7 +408,7 @@ namespace Lens { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/src/model/diffusion/lingbot_video.hpp b/src/model/diffusion/lingbot_video.hpp index 1868daa0..94af7199 100644 --- a/src/model/diffusion/lingbot_video.hpp +++ b/src/model/diffusion/lingbot_video.hpp @@ -674,7 +674,7 @@ namespace LingBotVideo { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/src/model/diffusion/ltxv.hpp b/src/model/diffusion/ltxv.hpp index b75c9f7f..e33448d2 100644 --- a/src/model/diffusion/ltxv.hpp +++ b/src/model/diffusion/ltxv.hpp @@ -1998,7 +1998,7 @@ namespace LTXV { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, audio_x, audio_timesteps, audio_length, frame_rate, video_positions); }; - auto out = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + auto out = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); return out; } diff --git a/src/model/diffusion/mage_flow.hpp b/src/model/diffusion/mage_flow.hpp index 6ac2d2aa..1531fb15 100644 --- a/src/model/diffusion/mage_flow.hpp +++ b/src/model/diffusion/mage_flow.hpp @@ -142,7 +142,7 @@ namespace MageFlow { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, ref_latents); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index a0a6c7ae..a927cd65 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -1168,8 +1168,6 @@ namespace MiniMaxH3 { }; return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, - false, - false, false), params.x->dim()); } diff --git a/src/model/diffusion/minit2i.hpp b/src/model/diffusion/minit2i.hpp index 28466105..110a7f76 100644 --- a/src/model/diffusion/minit2i.hpp +++ b/src/model/diffusion/minit2i.hpp @@ -589,7 +589,7 @@ namespace MiniT2I { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, mask); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/src/model/diffusion/mmdit.hpp b/src/model/diffusion/mmdit.hpp index 6731b5fb..2982742c 100644 --- a/src/model/diffusion/mmdit.hpp +++ b/src/model/diffusion/mmdit.hpp @@ -987,7 +987,7 @@ struct MMDiTRunner : public DiffusionModelRunner { return build_graph(x, timesteps, context, y, skip_layers); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/src/model/diffusion/pid.hpp b/src/model/diffusion/pid.hpp index 2a698e0a..24db7765 100644 --- a/src/model/diffusion/pid.hpp +++ b/src/model/diffusion/pid.hpp @@ -938,7 +938,7 @@ namespace Pid { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, lq_latent, degrade_sigma); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/src/model/diffusion/qwen_image.hpp b/src/model/diffusion/qwen_image.hpp index 6f08d618..14d33916 100644 --- a/src/model/diffusion/qwen_image.hpp +++ b/src/model/diffusion/qwen_image.hpp @@ -707,7 +707,7 @@ namespace Qwen { return build_graph(x, timesteps, context, ref_latents, ref_index_mode); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/src/model/diffusion/unet.hpp b/src/model/diffusion/unet.hpp index 56757304..3f2a87aa 100644 --- a/src/model/diffusion/unet.hpp +++ b/src/model/diffusion/unet.hpp @@ -835,7 +835,7 @@ struct UNetModelRunner : public DiffusionModelRunner { return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength, ip_context, ip_scale); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/src/model/diffusion/wan.hpp b/src/model/diffusion/wan.hpp index 9a907dcf..dc7042b1 100644 --- a/src/model/diffusion/wan.hpp +++ b/src/model/diffusion/wan.hpp @@ -950,7 +950,7 @@ namespace WAN { return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/src/model/diffusion/z_image.hpp b/src/model/diffusion/z_image.hpp index 176adfc1..d2604261 100644 --- a/src/model/diffusion/z_image.hpp +++ b/src/model/diffusion/z_image.hpp @@ -636,7 +636,7 @@ namespace ZImage { return build_graph(x, timesteps, context, ref_latents, ref_index_mode); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } sd::Tensor compute(int n_threads, diff --git a/src/model/te/clip.hpp b/src/model/te/clip.hpp index 2fde3de7..83d2646d 100644 --- a/src/model/te/clip.hpp +++ b/src/model/te/clip.hpp @@ -568,13 +568,11 @@ struct CLIPTextModelRunner : public GGMLRunner { size_t max_token_idx, bool return_pooled, int clip_skip, - bool auto_free = true, - bool free_compute_buffer = true, - bool free_compute_params = true) { + bool auto_runner_end = true) { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(input_ids, num_custom_embeddings, custom_embeddings_data, max_token_idx, return_pooled, clip_skip); }; - auto result = GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params); + auto result = GGMLRunner::compute(get_graph, n_threads, auto_runner_end); if (return_pooled) { return take_or_empty(std::move(result)); } diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp index f1a057d7..79752d4b 100644 --- a/src/model/te/llm.hpp +++ b/src/model/te/llm.hpp @@ -2079,9 +2079,7 @@ namespace LLM { const ImageEmbeds& image_embeds, std::set out_layers, bool return_all_hidden_states = false, - bool auto_free = true, - bool free_compute_buffer = true, - bool free_compute_params = true, + bool auto_runner_end = true, const DeepStackImageEmbeds& deepstack_image_embeds = {}, const std::vector& image_grids = {}) { auto get_graph = [&]() -> ggml_cgraph* { @@ -2093,7 +2091,7 @@ namespace LLM { out_layers, return_all_hidden_states); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params), + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, auto_runner_end), input_ids.dim() + 1); } @@ -2173,13 +2171,11 @@ namespace LLM { sd::Tensor encode_image(const int n_threads, const sd::Tensor& image, - bool auto_free = false, - bool free_compute_buffer = false, - bool free_compute_params = false) { + bool auto_runner_end = false) { auto get_graph = [&]() -> ggml_cgraph* { return build_encode_image_graph(image); }; - return take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params)); + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end)); } ggml_cgraph* build_encode_image_outputs_graph(const sd::Tensor& image_tensor) { @@ -2287,13 +2283,11 @@ namespace LLM { std::vector> encode_image_outputs(const int n_threads, const sd::Tensor& image, - bool auto_free = false, - bool free_compute_buffer = false, - bool free_compute_params = false) { + bool auto_runner_end = false) { auto get_graph = [&]() -> ggml_cgraph* { return build_encode_image_outputs_graph(image); }; - auto combined = take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params)); + auto combined = take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end)); if (combined.empty()) { return {}; } @@ -2312,20 +2306,14 @@ namespace LLM { std::vector> encode_video_block_outputs(const int n_threads, const sd::Tensor& frames, - bool auto_free = false, - bool free_compute_buffer = false, - bool free_compute_params = false) { + bool auto_runner_end = false) { int grid_h = static_cast(frames.shape()[1] / config.vision.patch_size); int grid_w = static_cast(frames.shape()[0] / config.vision.patch_size); auto pixel_values = process_video_block_tensor(frames, config.vision); auto get_graph = [&]() -> ggml_cgraph* { return build_encode_video_block_outputs_graph(pixel_values, grid_h, grid_w); }; - auto combined = take_or_empty(GGMLRunner::compute(get_graph, - n_threads, - auto_free, - free_compute_buffer, - free_compute_params)); + auto combined = take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end)); if (combined.empty()) { return {}; } diff --git a/src/model/te/t5.hpp b/src/model/te/t5.hpp index 90c781e1..f9f643f9 100644 --- a/src/model/te/t5.hpp +++ b/src/model/te/t5.hpp @@ -451,13 +451,11 @@ struct T5Runner : public GGMLRunner { sd::Tensor compute(const int n_threads, const sd::Tensor& input_ids, const sd::Tensor& attention_mask, - bool auto_free = true, - bool free_compute_buffer = true, - bool free_compute_params = true) { + bool auto_runner_end = true) { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(input_ids, attention_mask); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params), 3); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, auto_runner_end), 3); } static std::vector _relative_position_bucket(const std::vector& relative_position, diff --git a/src/model/upscaler/esrgan.hpp b/src/model/upscaler/esrgan.hpp index 21c97712..be02fb70 100644 --- a/src/model/upscaler/esrgan.hpp +++ b/src/model/upscaler/esrgan.hpp @@ -265,7 +265,7 @@ struct ESRGAN : public GGMLRunner { sd::Tensor compute(const int n_threads, const sd::Tensor& x) { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); }; - auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); return result; } }; diff --git a/src/model/upscaler/ltx_latent_upscaler.hpp b/src/model/upscaler/ltx_latent_upscaler.hpp index b70e1613..388799b0 100644 --- a/src/model/upscaler/ltx_latent_upscaler.hpp +++ b/src/model/upscaler/ltx_latent_upscaler.hpp @@ -499,7 +499,7 @@ namespace LTXVUpsampler { } size_t expected_dim = static_cast(x.dim()); auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), expected_dim); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), expected_dim); } }; diff --git a/src/model/vae/auto_encoder_kl.hpp b/src/model/vae/auto_encoder_kl.hpp index 604347d7..bf0a0381 100644 --- a/src/model/vae/auto_encoder_kl.hpp +++ b/src/model/vae/auto_encoder_kl.hpp @@ -744,7 +744,7 @@ struct AutoEncoderKL : public VAE { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(z, decode_graph); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), z.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), z.dim()); } sd::Tensor gaussian_latent_sample(const sd::Tensor& moments, std::shared_ptr rng) { diff --git a/src/model/vae/hunyuan_vae.hpp b/src/model/vae/hunyuan_vae.hpp index 938f5ffe..1f660364 100644 --- a/src/model/vae/hunyuan_vae.hpp +++ b/src/model/vae/hunyuan_vae.hpp @@ -827,9 +827,7 @@ namespace Hunyuan { }; auto output = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, - true, - true, - true), + false), graph_input.dim()); if (!output.empty() && input.dim() == 4) { output.squeeze_(2); diff --git a/src/model/vae/ltx_audio_vae.hpp b/src/model/vae/ltx_audio_vae.hpp index 3319a5c3..95ad01ce 100644 --- a/src/model/vae/ltx_audio_vae.hpp +++ b/src/model/vae/ltx_audio_vae.hpp @@ -1042,7 +1042,7 @@ namespace LTXV { ggml_build_forward_expand(gf, waveform); return gf; }; - auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), 4); + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), 4); int64_t t1 = ggml_time_ms(); LOG_INFO("ltx audio vae decode completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); return result; diff --git a/src/model/vae/ltx_vae.hpp b/src/model/vae/ltx_vae.hpp index 5629b939..53971d93 100644 --- a/src/model/vae/ltx_vae.hpp +++ b/src/model/vae/ltx_vae.hpp @@ -1334,7 +1334,6 @@ struct LTXVideoVAE : public VAE { (int)plan.tiles.size()); free_cache_ctx_and_buffer(); - cache_tensor_map.clear(); auto output = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor& z_chunk, const VAETemporalTile& tile) { LOG_DEBUG("LTX VAE temporal tile %lld/%d: latent frames [%lld, %lld), overlap=%d", @@ -1349,12 +1348,11 @@ struct LTXVideoVAE : public VAE { static_cast(tile.start), tile.overlap); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true, true, true), + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), expected_dim); }); free_cache_ctx_and_buffer(); - cache_tensor_map.clear(); return output; } @@ -1407,7 +1405,7 @@ struct LTXVideoVAE : public VAE { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(input, decode_graph); }; - auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), expected_dim); + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), expected_dim); if (result.empty()) { return {}; } @@ -1420,7 +1418,7 @@ struct LTXVideoVAE : public VAE { auto get_graph = [&]() -> ggml_cgraph* { return build_latent_statistics_graph(z, normalize); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), static_cast(z.dim())); } diff --git a/src/model/vae/mage_vae.hpp b/src/model/vae/mage_vae.hpp index 39075a83..71612afd 100644 --- a/src/model/vae/mage_vae.hpp +++ b/src/model/vae/mage_vae.hpp @@ -490,7 +490,7 @@ namespace MageVAE { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(input, decode_graph); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), input.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), input.dim()); } int get_encoder_output_channels(int input_channels) override { diff --git a/src/model/vae/minimax_h3_audio_vae.hpp b/src/model/vae/minimax_h3_audio_vae.hpp index 21f6ce16..6cd13daa 100644 --- a/src/model/vae/minimax_h3_audio_vae.hpp +++ b/src/model/vae/minimax_h3_audio_vae.hpp @@ -480,7 +480,7 @@ namespace MiniMaxH3 { return graph; }; auto result = restore_trailing_singleton_dims( - GGMLRunner::compute(get_graph, n_threads, false, false, false), + GGMLRunner::compute(get_graph, n_threads, false), 4); int64_t t1 = ggml_time_ms(); LOG_INFO("MiniMax-H3 audio VAE encode completed, taking %.2fs", @@ -500,7 +500,7 @@ namespace MiniMaxH3 { return graph; }; auto result = restore_trailing_singleton_dims( - GGMLRunner::compute(get_graph, n_threads, false, false, false), + GGMLRunner::compute(get_graph, n_threads, false), 4); int64_t t1 = ggml_time_ms(); LOG_INFO("MiniMax-H3 audio VAE decode completed, taking %.2fs", diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index a8aa6b5f..dcfea988 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -793,8 +793,6 @@ namespace MiniMaxH3VAE { return restore_trailing_singleton_dims( GGMLRunner::compute(get_graph, n_threads, - false, - false, false), 5); } diff --git a/src/model/vae/tae.hpp b/src/model/vae/tae.hpp index d291bb78..caa580fc 100644 --- a/src/model/vae/tae.hpp +++ b/src/model/vae/tae.hpp @@ -787,7 +787,7 @@ struct TinyImageAutoEncoder : public VAE { return build_graph(z_tensor, decode_graph); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), z_tensor.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), z_tensor.dim()); } }; @@ -872,7 +872,7 @@ struct TinyVideoAutoEncoder : public VAE { return build_graph(z_tensor, decode_graph); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), z_tensor.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), z_tensor.dim()); } }; diff --git a/src/model/vae/vae.hpp b/src/model/vae/vae.hpp index f3adb0cc..758fb5ee 100644 --- a/src/model/vae/vae.hpp +++ b/src/model/vae/vae.hpp @@ -251,7 +251,7 @@ public: tiling_params); } - runner_done(); + runner_end(); if (output.empty()) { LOG_ERROR("vae encode compute failed"); @@ -305,7 +305,7 @@ public: tiling_params); } - runner_done(); + runner_end(); if (output.empty()) { LOG_ERROR("vae decode compute failed"); diff --git a/src/model/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp index 422a5782..517b49c9 100644 --- a/src/model/vae/wan_vae.hpp +++ b/src/model/vae/wan_vae.hpp @@ -1415,7 +1415,6 @@ namespace WAN { (int)plan.tiles.size()); free_cache_ctx_and_buffer(); - cache_tensor_map.clear(); ae.clear_cache(); auto output = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor& input_tile, const VAETemporalTile& tile) { @@ -1428,12 +1427,11 @@ namespace WAN { return build_temporal_tile_graph(input_tile, static_cast(tile.start)); }; return restore_trailing_singleton_dims( - GGMLRunner::compute(get_graph, n_threads, true, true, true), + GGMLRunner::compute(get_graph, n_threads, false), static_cast(input.dim())); }); free_cache_ctx_and_buffer(); - cache_tensor_map.clear(); ae.clear_cache(); return output; } @@ -1448,7 +1446,7 @@ namespace WAN { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(input.empty() ? z : input, decode_graph); }; - auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true, true, true), + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), input.empty() ? z.dim() : input.dim()); if (!result.empty() && z.dim() == 4) { result.squeeze_(2); diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 72e81374..e7c52af4 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -143,7 +143,7 @@ size_t estimate_tensors_size(const std::map& tensors) return size; } -void ModelManager::set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft) { +void ModelManager::set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft, const std::vector>& device_limits) { if (compute_backend == nullptr) { return; } @@ -152,6 +152,7 @@ void ModelManager::set_split_buffer_type(ggml_backend_t compute_backend, ggml_ba return; } split_buffer_types_[compute_backend] = split_buft; + split_buffer_devices_[split_buft] = device_limits; } bool ModelManager::tensor_shape_supports_split_buffer(const ggml_tensor* tensor) { @@ -164,11 +165,10 @@ bool ModelManager::tensor_shape_supports_split_buffer(const ggml_tensor* tensor) } ggml_backend_buffer_type_t ModelManager::split_buffer_type_for(const TensorState& state) const { - if (!state.allow_split_buffer || !tensor_shape_supports_split_buffer(state.tensor)) { + if (!tensor_shape_supports_split_buffer(state.tensor)) { return nullptr; } - auto it = split_buffer_types_.find(state.compute_backend); - return it != split_buffer_types_.end() ? it->second : nullptr; + return state.split_buffer_type; } bool ModelManager::register_param_tensors(const std::string& desc, @@ -203,14 +203,17 @@ bool ModelManager::register_param_tensors(const std::string& desc, } ggml_set_name(tensor, name.c_str()); - auto state = std::make_unique(); - state->name = name; - state->tensor = tensor; - state->desc = desc; - state->residency_mode = residency_mode; - state->compute_backend = compute_backend; - state->params_backend = params_backend; - state->allow_split_buffer = allow_split_buffer; + auto state = std::make_unique(); + state->name = name; + state->tensor = tensor; + state->desc = desc; + state->residency_mode = residency_mode; + state->compute_backend = compute_backend; + state->params_backend = params_backend; + auto split_buffer = split_buffer_types_.find(compute_backend); + if (allow_split_buffer && split_buffer != split_buffer_types_.end()) { + state->split_buffer_type = split_buffer->second; + } state->params_follow_compute_backend = params_follow_compute_backend; if (tensor_ops != nullptr) { auto op_it = tensor_ops->find(tensor); @@ -240,7 +243,7 @@ bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* reg if (state == nullptr || state->desc != desc) { continue; } - if (state->active_prepare_count > 0) { + if (state->pin_count > 0) { LOG_ERROR("model manager cannot unregister active %s tensor '%s'", desc.c_str(), state->name.c_str()); @@ -287,7 +290,7 @@ bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* reg if (state == nullptr) { continue; } - if (state->active_prepare_count > 0 || state->staged_to_compute_backend) { + if (state->pin_count > 0 || state->staged_to_compute_backend) { LOG_ERROR("model manager cannot unregister %s while tensor '%s' is active", desc.c_str(), state->name.c_str()); @@ -403,14 +406,27 @@ bool ModelManager::load_tensors_to_params_backend(const std::vector prepared; for (ParamsStorageBlock* block : created_storage_blocks) { if (block != nullptr && block->buffer != nullptr) { - LOG_DEBUG("model manager prepared params backend buffer (%6.2f MB, %zu tensors, %s)", - ggml_backend_buffer_get_size(block->buffer) / (1024.f * 1024.f), - block->states.size(), - ggml_backend_buffer_is_host(block->buffer) ? "RAM" : "VRAM"); + auto& stats = prepared[ggml_backend_buffer_get_type(block->buffer)]; + stats.bytes += ggml_backend_buffer_get_size(block->buffer); + stats.tensors += block->states.size(); + ++stats.blocks; } } + for (const auto& entry : prepared) { + LOG_DEBUG("model manager prepared params backend buffers (%6.2f MB, %zu tensors, %zu blocks, %s) on %s", + entry.second.bytes / (1024.f * 1024.f), + entry.second.tensors, entry.second.blocks, + ggml_backend_buft_is_host(entry.first) ? "RAM" : "VRAM", + ggml_backend_buft_name(entry.first)); + } return true; } @@ -444,68 +460,99 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector& states = pair.second; - if (states.empty()) { + ggml_backend_t compute_backend = pair.first.first; + ggml_backend_buffer_type_t staging_buft = pair.first.second; + const std::vector& target_states = pair.second; + if (target_states.empty()) { continue; } - int64_t t0 = ggml_time_ms(); - - ggml_init_params init_params; - init_params.mem_size = std::max(1, states.size()) * ggml_tensor_overhead(); - init_params.mem_buffer = nullptr; - init_params.no_alloc = true; - - ggml_context* staging_ctx = ggml_init(init_params); - GGML_ASSERT(staging_ctx != nullptr); - - std::vector> staged_tensors; - staged_tensors.reserve(states.size()); - for (TensorState* state : states) { - ggml_tensor* staging_tensor = ggml_dup_tensor(staging_ctx, state->tensor); - ggml_set_name(staging_tensor, state->tensor->name); - staged_tensors.push_back({state, staging_tensor}); + const size_t alignment = ggml_backend_buft_get_alignment(staging_buft); + size_t backend_limit = ggml_backend_buft_get_max_size(staging_buft); + if (!ggml_backend_buft_is_host(staging_buft) && + (backend_limit == 0 || backend_limit > MAX_RESIDENCY_BLOCK_BYTES)) { + backend_limit = MAX_RESIDENCY_BLOCK_BYTES; } - ggml_backend_buffer_t compute_buffer = ggml_backend_alloc_ctx_tensors_from_buft(staging_ctx, staging_buft); - if (compute_buffer == nullptr) { - LOG_ERROR("model manager alloc compute params backend buffer failed, num_tensors = %zu", - staged_tensors.size()); - ggml_free(staging_ctx); + const int64_t t0 = ggml_time_ms(); + size_t staged_bytes = 0; + size_t staged_blocks = 0; + auto stage_chunk = [&](const std::vector& chunk) -> bool { + if (chunk.empty()) { + return true; + } + ggml_init_params init_params; + init_params.mem_size = std::max(1, chunk.size()) * ggml_tensor_overhead(); + init_params.mem_buffer = nullptr; + init_params.no_alloc = true; + + ggml_context* staging_ctx = ggml_init(init_params); + GGML_ASSERT(staging_ctx != nullptr); + std::vector> staged_tensors; + staged_tensors.reserve(chunk.size()); + for (TensorState* state : chunk) { + ggml_tensor* staging_tensor = ggml_dup_tensor(staging_ctx, state->tensor); + ggml_set_name(staging_tensor, state->tensor->name); + staged_tensors.push_back({state, staging_tensor}); + } + + ggml_backend_buffer_t compute_buffer = + ggml_backend_alloc_ctx_tensors_from_buft(staging_ctx, staging_buft); + if (compute_buffer == nullptr) { + LOG_ERROR("model manager alloc compute params backend buffer failed, num_tensors = %zu", + staged_tensors.size()); + ggml_free(staging_ctx); + return false; + } + ggml_backend_buffer_set_usage(compute_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + for (auto& staged_tensor : staged_tensors) { + TensorState* state = staged_tensor.first; + ggml_tensor* managed_tensor = state->tensor; + ggml_tensor* staging_tensor = staged_tensor.second; + ggml_backend_tensor_copy(managed_tensor, staging_tensor); + std::swap(managed_tensor->buffer, staging_tensor->buffer); + std::swap(managed_tensor->data, staging_tensor->data); + std::swap(managed_tensor->extra, staging_tensor->extra); + state->staged_to_compute_backend = true; + } + ggml_backend_synchronize(compute_backend); + + auto block = std::make_unique(); + block->compute_backend = compute_backend; + block->buffer = compute_buffer; + block->staging_ctx = staging_ctx; + block->staged_tensors = std::move(staged_tensors); + staged_bytes += ggml_backend_buffer_get_size(compute_buffer); + ++staged_blocks; + compute_staging_blocks_.push_back(std::move(block)); + return true; + }; + + std::vector chunk; + size_t chunk_size = 0; + for (TensorState* state : target_states) { + const size_t tensor_size = GGML_PAD( + ggml_backend_buft_get_alloc_size(staging_buft, state->tensor), alignment); + if (!chunk.empty() && backend_limit > 0 && + tensor_size > backend_limit - std::min(chunk_size, backend_limit)) { + if (!stage_chunk(chunk)) { + return false; + } + chunk.clear(); + chunk_size = 0; + } + chunk.push_back(state); + chunk_size = tensor_size > SIZE_MAX - chunk_size ? SIZE_MAX : chunk_size + tensor_size; + } + if (!stage_chunk(chunk)) { return false; } - ggml_backend_buffer_set_usage(compute_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - - for (auto& staged_tensor : staged_tensors) { - TensorState* state = staged_tensor.first; - ggml_tensor* managed_tensor = state->tensor; - ggml_tensor* staging_tensor = staged_tensor.second; - ggml_backend_tensor_copy(managed_tensor, staging_tensor); - std::swap(managed_tensor->buffer, staging_tensor->buffer); - std::swap(managed_tensor->data, staging_tensor->data); - std::swap(managed_tensor->extra, staging_tensor->extra); - } - ggml_backend_synchronize(compute_backend); - - auto block = std::make_unique(); - block->compute_backend = compute_backend; - block->buffer = compute_buffer; - block->staging_ctx = staging_ctx; - block->staged_tensors = std::move(staged_tensors); - for (auto& staged_tensor : block->staged_tensors) { - TensorState* state = staged_tensor.first; - state->staged_to_compute_backend = true; - } - compute_staging_blocks_.push_back(std::move(block)); - - int64_t t1 = ggml_time_ms(); - LOG_DEBUG("model manager staged compute params (%6.2f MB, %zu tensors) to %s, taking %.2fs", - ggml_backend_buffer_get_size(compute_buffer) / (1024.f * 1024.f), - states.size(), + LOG_DEBUG("model manager staged compute params (%6.2f MB, %zu tensors, %zu blocks) to %s, taking %.2fs", + staged_bytes / (1024.f * 1024.f), + target_states.size(), + staged_blocks, ggml_backend_name(compute_backend), - (t1 - t0) * 1.0f / 1000); + (ggml_time_ms() - t0) / 1000.f); } return true; @@ -729,6 +776,10 @@ bool ModelManager::alloc_params_buffers(const std::vector& states, const std::vector& states = pair.second; size_t alignment = ggml_backend_buft_get_alignment(params_buft); size_t max_size = ggml_backend_buft_get_max_size(params_buft); + if (!ggml_backend_buft_is_host(params_buft) && + (max_size == 0 || max_size > MAX_RESIDENCY_BLOCK_BYTES)) { + max_size = MAX_RESIDENCY_BLOCK_BYTES; + } auto alloc_chunk = [&](const std::vector& chunk, size_t chunk_size) -> bool { if (chunk.empty() || chunk_size == 0) { @@ -935,10 +986,6 @@ void ModelManager::free_compute_staging_block(ComputeStagingBlock& block) { } if (block.buffer != nullptr) { - LOG_DEBUG("model manager releasing compute params (%6.2f MB, %zu tensors) from %s", - ggml_backend_buffer_get_size(block.buffer) / (1024.f * 1024.f), - block.staged_tensors.size(), - block.compute_backend != nullptr ? ggml_backend_name(block.compute_backend) : "unknown"); ggml_backend_buffer_free(block.buffer); block.buffer = nullptr; } @@ -951,6 +998,12 @@ void ModelManager::free_compute_staging_block(ComputeStagingBlock& block) { void ModelManager::release_compute_staging_blocks(bool force, const std::unordered_set* target_states) { + struct ReleaseStats { + size_t bytes = 0; + size_t tensors = 0; + size_t blocks = 0; + }; + std::map released; for (auto it = compute_staging_blocks_.begin(); it != compute_staging_blocks_.end();) { ComputeStagingBlock* block = it->get(); bool can_release = force; @@ -966,25 +1019,33 @@ void ModelManager::release_compute_staging_blocks(bool force, target_states->find(state) == target_states->end()) { return false; } - return state->active_prepare_count == 0; + return state->pin_count == 0; }); } if (can_release) { + if (block->buffer != nullptr) { + auto& stats = released[block->compute_backend]; + stats.bytes += ggml_backend_buffer_get_size(block->buffer); + stats.tensors += block->staged_tensors.size(); + ++stats.blocks; + } free_compute_staging_block(*block); it = compute_staging_blocks_.erase(it); } else { ++it; } } + for (const auto& entry : released) { + LOG_DEBUG("model manager released compute params (%6.2f MB, %zu tensors, %zu blocks) from %s", + entry.second.bytes / (1024.f * 1024.f), + entry.second.tensors, entry.second.blocks, + entry.first != nullptr ? ggml_backend_name(entry.first) : "unknown"); + } } void ModelManager::free_params_storage_block(ParamsStorageBlock& block) { if (block.buffer != nullptr) { - LOG_DEBUG("model manager releasing params backend buffer (%6.2f MB, %zu tensors, %s)", - ggml_backend_buffer_get_size(block.buffer) / (1024.f * 1024.f), - block.states.size(), - ggml_backend_buffer_is_host(block.buffer) ? "RAM" : "VRAM"); ggml_backend_buffer_free(block.buffer); block.buffer = nullptr; } @@ -1006,6 +1067,12 @@ void ModelManager::free_params_storage_block(ParamsStorageBlock& block) { void ModelManager::release_params_storage_blocks(bool force, const std::unordered_set* target_states) { + struct ReleaseStats { + size_t bytes = 0; + size_t tensors = 0; + size_t blocks = 0; + }; + std::map released; for (auto it = params_storage_blocks_.begin(); it != params_storage_blocks_.end();) { ParamsStorageBlock* block = it->get(); bool can_release = force; @@ -1020,19 +1087,32 @@ void ModelManager::release_params_storage_blocks(bool force, target_states->find(state) == target_states->end()) { return false; } - return state->active_prepare_count == 0 && + return state->pin_count == 0 && !state->staged_to_compute_backend && state->residency_mode == ResidencyMode::Disk; }); } if (can_release) { + if (block->buffer != nullptr) { + auto& stats = released[ggml_backend_buffer_get_type(block->buffer)]; + stats.bytes += ggml_backend_buffer_get_size(block->buffer); + stats.tensors += block->states.size(); + ++stats.blocks; + } free_params_storage_block(*block); it = params_storage_blocks_.erase(it); } else { ++it; } } + for (const auto& entry : released) { + LOG_DEBUG("model manager released params backend buffers (%6.2f MB, %zu tensors, %zu blocks, %s) from %s", + entry.second.bytes / (1024.f * 1024.f), + entry.second.tensors, entry.second.blocks, + ggml_backend_buft_is_host(entry.first) ? "RAM" : "VRAM", + ggml_backend_buft_name(entry.first)); + } } void ModelManager::erase_params_storage_block(ParamsStorageBlock* block) { @@ -1048,16 +1128,19 @@ void ModelManager::erase_params_storage_block(ParamsStorageBlock* block) { void ModelManager::release_all() { clear_all_prefetched_params(); + runtime_residencies_.clear(); + workspace_reclaimers_.clear(); for (auto& state : tensor_states_) { - state->active_prepare_count = 0; - state->applied_lora_epoch = UINT64_MAX; + state->pin_count = 0; + state->applied_lora_epoch = UINT64_MAX; } release_compute_staging_blocks(true); release_params_storage_blocks(true); } bool ModelManager::resolve_required_tensor_states(const std::vector& tensors, - std::vector& required_states) const { + std::vector& required_states, + ggml_backend_t compute_backend) const { required_states.clear(); std::unordered_set seen; for (ggml_tensor* tensor : tensors) { @@ -1079,7 +1162,9 @@ bool ModelManager::resolve_required_tensor_states(const std::vectorcompute_backend == nullptr || + state->compute_backend == compute_backend) && + seen.insert(state).second) { required_states.push_back(state); } } @@ -1115,7 +1200,7 @@ bool ModelManager::assign_compute_backend(const std::vector& tenso continue; } - if (state->active_prepare_count > 0 || state->staged_to_compute_backend) { + if (state->pin_count > 0 || state->staged_to_compute_backend) { LOG_ERROR("model manager cannot move active tensor '%s' to another compute backend", state->name.c_str()); return false; @@ -1135,6 +1220,131 @@ bool ModelManager::assign_compute_backend(const std::vector& tenso return true; } +size_t ModelManager::compute_backend_alloc_size(const std::vector& states, + bool missing_only) const { + size_t total_size = 0; + std::unordered_set seen; + for (TensorState* state : states) { + if (state == nullptr || state->tensor == nullptr || !seen.insert(state).second || + should_ignore(*state) || is_optional_missing_tensor(state->name)) { + continue; + } + const bool compute_resident = + state->compute_backend == state->params_backend + ? state->loaded_to_params_backend + : state->staged_to_compute_backend; + if (missing_only && compute_resident) { + continue; + } + + ggml_backend_buffer_type_t buffer_type = nullptr; + if (state->compute_backend == state->params_backend) { + buffer_type = params_buffer_type_for(*state); + } else { + buffer_type = split_buffer_type_for(*state); + if (buffer_type == nullptr && state->compute_backend != nullptr) { + buffer_type = ggml_backend_get_default_buffer_type(state->compute_backend); + } + } + if (buffer_type == nullptr) { + continue; + } + const size_t alignment = ggml_backend_buft_get_alignment(buffer_type); + const size_t tensor_size = ggml_backend_buft_get_alloc_size(buffer_type, state->tensor); + const size_t alloc_size = GGML_PAD(tensor_size, alignment); + if (alloc_size > SIZE_MAX - total_size) { + return SIZE_MAX; + } + total_size += alloc_size; + } + return total_size; +} + +size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backend) const { + if (compute_backend == nullptr) { + return 0; + } + ggml_backend_dev_t compute_device = ggml_backend_get_device(compute_backend); + if (compute_device == nullptr) { + return 0; + } + + size_t total_size = 0; + auto add_buffer = [&](ggml_backend_buffer_t buffer) { + if (buffer == nullptr || ggml_backend_buffer_is_host(buffer)) { + return; + } + ggml_backend_buffer_type_t buffer_type = ggml_backend_buffer_get_type(buffer); + auto split_devices = split_buffer_devices_.find(buffer_type); + const bool on_device = split_devices == split_buffer_devices_.end() + ? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device + : std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) { + return ggml_backend_get_device(entry.first) == compute_device; + }); + if (!on_device) { + return; + } + const size_t buffer_size = ggml_backend_buffer_get_size(buffer); + total_size = buffer_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + buffer_size; + }; + + for (const auto& block : params_storage_blocks_) { + if (block != nullptr) { + add_buffer(block->buffer); + } + } + for (const auto& block : compute_staging_blocks_) { + if (block != nullptr) { + add_buffer(block->buffer); + } + } + for (const auto& entry : prefetch_blocks_) { + if (entry.second != nullptr) { + for (const auto& block : entry.second->staging_blocks) { + if (block != nullptr) { + add_buffer(block->buffer); + } + } + } + } + return total_size; +} + +void ModelManager::update_runtime_residency(uintptr_t owner_id, + ggml_backend_t compute_backend, + size_t resident_bytes) { + if (owner_id == 0) { + return; + } + if (compute_backend == nullptr || resident_bytes == 0) { + runtime_residencies_.erase({owner_id, compute_backend}); + return; + } + runtime_residencies_[{owner_id, compute_backend}] = {compute_backend, resident_bytes}; +} + +size_t ModelManager::other_runtime_resident_bytes(uintptr_t owner_id, + ggml_backend_t compute_backend) const { + if (compute_backend == nullptr) { + return 0; + } + ggml_backend_dev_t compute_device = ggml_backend_get_device(compute_backend); + if (compute_device == nullptr) { + return 0; + } + size_t total_size = 0; + for (const auto& entry : runtime_residencies_) { + if (entry.first.first == owner_id || entry.second.compute_backend == nullptr || + ggml_backend_get_device(entry.second.compute_backend) != compute_device) { + continue; + } + total_size = entry.second.resident_bytes > SIZE_MAX - total_size + ? SIZE_MAX + : total_size + entry.second.resident_bytes; + } + return total_size; +} + bool ModelManager::prepare_params(const std::vector& tensors) { if (tensors.empty()) { return true; @@ -1155,18 +1365,20 @@ bool ModelManager::prepare_params(const std::vector& tensors) { return false; } + // LoRA execution may reclaim other residency blocks while these weights are in use. + const uint64_t use_epoch = ++residency_epoch_; + for (TensorState* state : required_states) { + if (state != nullptr) { + state->pin_count++; + state->last_use_epoch = use_epoch; + } + } if (!apply_loras_to_params(required_states)) { + finish_compute_backend_usage(required_states); release_compute_staging_blocks(false); release_params_storage_blocks(false); return false; } - - for (TensorState* state : required_states) { - if (state == nullptr) { - continue; - } - state->active_prepare_count++; - } return true; } @@ -1180,11 +1392,10 @@ void ModelManager::finish_compute_backend_usage(const std::vector& if (state == nullptr || !target_states.insert(state).second) { continue; } - if (state->active_prepare_count > 0) { - state->active_prepare_count--; + if (state->pin_count > 0) { + state->pin_count--; } } - release_compute_staging_blocks(false, &target_states); } void ModelManager::release_compute_backend_params(const std::vector& tensors) { @@ -1198,7 +1409,7 @@ void ModelManager::release_compute_backend_params(const std::vector& tensors) { +void ModelManager::evict_compute_backend_params(const std::vector& tensors) { if (tensors.empty()) { return; } @@ -1206,9 +1417,285 @@ void ModelManager::release_params_backend_params(const std::vector if (!resolve_required_tensor_states(tensors, required_states)) { return; } - if (required_states.empty()) { - return; - } std::unordered_set target_states(required_states.begin(), required_states.end()); + + for (const auto& block : compute_staging_blocks_) { + const bool intersects = std::any_of( + block->staged_tensors.begin(), + block->staged_tensors.end(), + [&](const std::pair& pair) { + return pair.first != nullptr && target_states.count(pair.first) > 0; + }); + const bool fully_evictable = std::all_of( + block->staged_tensors.begin(), + block->staged_tensors.end(), + [](const std::pair& pair) { + return pair.first == nullptr || pair.first->pin_count == 0; + }); + if (intersects && fully_evictable) { + for (const auto& pair : block->staged_tensors) { + if (pair.first != nullptr) { + target_states.insert(pair.first); + } + } + } + } + release_compute_staging_blocks(false, &target_states); + + for (const auto& block : params_storage_blocks_) { + const bool intersects = std::any_of( + block->states.begin(), + block->states.end(), + [&](TensorState* state) { + return state != nullptr && target_states.count(state) > 0; + }); + const bool fully_evictable = std::all_of( + block->states.begin(), + block->states.end(), + [](TensorState* state) { + return state == nullptr || + (state->pin_count == 0 && !state->staged_to_compute_backend && + state->residency_mode == ResidencyMode::Disk); + }); + if (intersects && fully_evictable) { + target_states.insert(block->states.begin(), block->states.end()); + } + } release_params_storage_blocks(false, &target_states); } +WeightResidencyInfo ModelManager::inspect_compute_backend_params( + const std::vector& tensors) const { + WeightResidencyInfo info; + std::vector states; + if (!resolve_required_tensor_states(tensors, states)) { + return info; + } + + ggml_backend_t prefetch_compute_backend = nullptr; + bool has_missing_params = false; + bool prefetch_candidate = true; + for (TensorState* state : states) { + if (state == nullptr || should_ignore(*state) || + is_optional_missing_tensor(state->name)) { + continue; + } + const bool compute_resident = + state->compute_backend == state->params_backend + ? state->loaded_to_params_backend + : state->staged_to_compute_backend; + if (compute_resident) { + continue; + } + has_missing_params = true; + if (split_buffer_type_for(*state) != nullptr) { + prefetch_candidate = false; + } + if (state->compute_backend == state->params_backend || + state->compute_backend == nullptr || sd_backend_is_cpu(state->compute_backend)) { + prefetch_candidate = false; + continue; + } + if (prefetch_compute_backend == nullptr) { + prefetch_compute_backend = state->compute_backend; + } else if (prefetch_compute_backend != state->compute_backend) { + prefetch_candidate = false; + } + } + info.missing_bytes = compute_backend_alloc_size(states, true); + if (has_missing_params && prefetch_candidate && prefetch_compute_backend != nullptr) { + ggml_backend_dev_t device = ggml_backend_get_device(prefetch_compute_backend); + if (device != nullptr) { + ggml_backend_dev_props props{}; + ggml_backend_dev_get_props(device, &props); + info.async_prefetch_supported = props.caps.async; + } + } + return info; +} + +void ModelManager::set_workspace_reclaimer(uintptr_t owner_id, std::function reclaim) { + workspace_reclaimers_[owner_id] = std::move(reclaim); +} + +void ModelManager::remove_runtime_owner(uintptr_t owner_id) { + workspace_reclaimers_.erase(owner_id); + for (auto it = runtime_residencies_.begin(); it != runtime_residencies_.end();) { + if (it->first.first == owner_id) { + it = runtime_residencies_.erase(it); + } else { + ++it; + } + } +} + +ModelManager::CapacityCheck ModelManager::check_capacity( + const DeviceMemoryRequest& request, + const std::vector& states) const { + CapacityCheck result; + 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) { + 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 (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)); + result.available_budget_bytes = resident < request.max_backend_bytes + ? request.max_backend_bytes - resident + : 0; + } + std::map split_devices; + for (auto state : states) { + auto placement = split_buffer_devices_.find(split_buffer_type_for(*state)); + if (placement != split_buffer_devices_.end()) { + for (const auto& entry : placement->second) { + auto inserted = split_devices.emplace(entry); + if (!inserted.second && entry.second > 0) { + auto& limit = inserted.first->second; + limit = limit == 0 ? entry.second : std::min(limit, entry.second); + } + } + } + } + // 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); + } + if (entry.second > 0) { + const size_t resident = add(compute_backend_resident_bytes(entry.first), + other_runtime_resident_bytes(request.owner_id, entry.first)); + const size_t available = resident < entry.second ? entry.second - resident : 0; + result.available_budget_bytes = std::min(result.available_budget_bytes, available); + } + } + return result; +} + +bool ModelManager::fits_compute_backend_capacity( + const DeviceMemoryRequest& request, + const std::vector& required_params) const { + std::vector states; + return resolve_required_tensor_states(required_params, states, request.compute_backend) && + check_capacity(request, states).fits(); +} + +bool ModelManager::ensure_compute_backend_capacity( + const DeviceMemoryRequest& request, + const std::vector& required_params, + const std::vector>& preferred_eviction_order, + const std::vector& protected_params) { + std::vector required_states; + if (!resolve_required_tensor_states(required_params, required_states, request.compute_backend)) { + return false; + } + + ggml_backend_t compute_backend = request.compute_backend; + if (compute_backend == nullptr) { + LOG_ERROR("model manager cannot reclaim memory for a null compute backend"); + return false; + } + if (sd_backend_is_cpu(compute_backend)) { + return true; + } + + auto fits = [&]() { return check_capacity(request, required_states).fits(); }; + if (fits()) { + return true; + } + for (const auto& entry : workspace_reclaimers_) { + if (entry.first != request.owner_id) { + entry.second(); + if (fits()) { + return true; + } + } + } + + std::unordered_set protected_states; + std::vector resolved_protected; + if (!resolve_required_tensor_states(protected_params, resolved_protected)) { + return false; + } + protected_states.insert(resolved_protected.begin(), resolved_protected.end()); + for (const auto& entry : prefetch_blocks_) { + if (entry.second != nullptr) { + protected_states.insert(entry.second->states.begin(), entry.second->states.end()); + } + } + + std::unordered_set eviction_states; + auto add_evictable_state = [&](TensorState* state) { + if (state == nullptr || state->compute_backend != compute_backend || + state->pin_count > 0 || protected_states.find(state) != protected_states.end()) { + return; + } + const bool reloadable = state->residency_mode == ResidencyMode::Disk || + state->compute_backend != state->params_backend; + const bool resident = state->compute_backend == state->params_backend + ? state->loaded_to_params_backend + : state->staged_to_compute_backend; + if (reloadable && resident) { + eviction_states.insert(state); + } + }; + auto release_eviction_states = [&]() { + release_compute_staging_blocks(false, &eviction_states); + release_params_storage_blocks(false, &eviction_states); + return fits(); + }; + + for (const auto& candidate_params : preferred_eviction_order) { + std::vector candidate_states; + if (!resolve_required_tensor_states(candidate_params, candidate_states)) { + return false; + } + for (TensorState* state : candidate_states) { + add_evictable_state(state); + } + if (release_eviction_states()) { + return true; + } + } + + std::vector global_candidates; + global_candidates.reserve(tensor_states_.size()); + for (const auto& state : tensor_states_) { + if (state != nullptr && eviction_states.find(state.get()) == eviction_states.end()) { + global_candidates.push_back(state.get()); + } + } + std::stable_sort(global_candidates.begin(), + global_candidates.end(), + [](const TensorState* lhs, const TensorState* rhs) { + return lhs->last_use_epoch < rhs->last_use_epoch; + }); + for (TensorState* state : global_candidates) { + add_evictable_state(state); + if (release_eviction_states()) { + return true; + } + } + + 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", + 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)); + return false; +} diff --git a/src/model_manager.h b/src/model_manager.h index d6910d1d..b23f0db6 100644 --- a/src/model_manager.h +++ b/src/model_manager.h @@ -9,10 +9,10 @@ #include #include +#include "device_residency_manager.h" #include "model_loader.h" -#include "weight_manager.h" -class ModelManager : public RunnerWeightManager { +class ModelManager : public DeviceResidencyManager { public: enum class ResidencyMode { Disk, @@ -28,24 +28,27 @@ public: }; private: + static constexpr size_t MAX_RESIDENCY_BLOCK_BYTES = 64ULL * 1024ULL * 1024ULL; + struct TensorState { std::string name; ggml_tensor* tensor = nullptr; std::string desc; - ResidencyMode residency_mode = ResidencyMode::ParamBackend; - ggml_backend_t compute_backend = nullptr; - ggml_backend_t params_backend = nullptr; - bool allow_split_buffer = false; - bool params_follow_compute_backend = false; - bool metadata_validated = false; - enum ggml_op usage_op = GGML_OP_NONE; + ResidencyMode residency_mode = ResidencyMode::ParamBackend; + ggml_backend_t compute_backend = nullptr; + ggml_backend_t params_backend = nullptr; + ggml_backend_buffer_type_t split_buffer_type = nullptr; + bool params_follow_compute_backend = false; + bool metadata_validated = false; + enum ggml_op usage_op = GGML_OP_NONE; - int active_prepare_count = 0; + int pin_count = 0; bool loaded_to_params_backend = false; bool staged_to_compute_backend = false; uint64_t applied_lora_epoch = UINT64_MAX; + uint64_t last_use_epoch = 0; }; struct ParamsStorageBlock { @@ -66,9 +69,12 @@ private: ggml_backend_t compute_backend = nullptr; ggml_backend_t transfer_backend = nullptr; ggml_backend_event_t event = nullptr; - ggml_context* staging_ctx = nullptr; - ggml_backend_buffer_t buffer = nullptr; - std::vector> staged_tensors; + std::vector> staging_blocks; + }; + + struct RuntimeResidency { + ggml_backend_t compute_backend = nullptr; + size_t resident_bytes = 0; }; ModelLoader model_loader_; @@ -77,16 +83,22 @@ private: std::vector> params_storage_blocks_; std::vector> compute_staging_blocks_; std::map split_buffer_types_; + std::map>> split_buffer_devices_; std::map> prefetch_blocks_; std::map prefetch_backends_; + std::map, RuntimeResidency> runtime_residencies_; + std::map> workspace_reclaimers_; bool warned_split_lora_skip_ = false; std::set common_ignore_tensors_; std::vector loras_; - SDVersion lora_version_ = VERSION_COUNT; - uint64_t current_lora_epoch_ = 0; - int n_threads_ = 0; - bool enable_mmap_ = false; - bool writable_mmap_ = false; + SDVersion lora_version_ = VERSION_COUNT; + uint64_t current_lora_epoch_ = 0; + uint64_t residency_epoch_ = 0; + int n_threads_ = 0; + bool enable_mmap_ = false; + bool writable_mmap_ = false; + bool segmented_compute_disabled_ = false; + bool prefetch_disabled_ = false; void finish_compute_backend_usage(const std::vector& states); void release_all(); @@ -99,7 +111,8 @@ private: void release_prefetch(); bool resolve_required_tensor_states(const std::vector& tensors, - std::vector& required_states) const; + std::vector& required_states, + ggml_backend_t compute_backend = nullptr) const; bool should_ignore(const TensorState& state) const; bool is_optional_missing_tensor(const std::string& name) const; bool validate_tensor(const TensorState& state) const; @@ -113,6 +126,21 @@ private: std::vector& created_storage_blocks); bool load_tensors(const std::vector& states); bool stage_tensors_to_compute_backend(const std::vector& states); + size_t compute_backend_alloc_size(const std::vector& states, + bool missing_only) const; + size_t compute_backend_resident_bytes(ggml_backend_t compute_backend) const; + struct CapacityCheck { + size_t required_device_bytes = 0; + size_t required_budget_bytes = 0; + size_t available_device_bytes = SIZE_MAX; + size_t available_budget_bytes = SIZE_MAX; + bool fits() const { + return required_device_bytes <= available_device_bytes && + required_budget_bytes <= available_budget_bytes; + } + }; + CapacityCheck check_capacity(const DeviceMemoryRequest& request, + const std::vector& states) const; ggml_backend_buffer_type_t params_buffer_type_for(const TensorState& state) const; ggml_backend_buffer_type_t split_buffer_type_for(const TensorState& state) const; @@ -124,6 +152,8 @@ private: void free_params_storage_block(ParamsStorageBlock& block); void erase_params_storage_block(ParamsStorageBlock* block); void reset_lora_applied_params(); + size_t other_runtime_resident_bytes(uintptr_t owner_id, + ggml_backend_t compute_backend) const; public: ~ModelManager() override; @@ -135,11 +165,15 @@ public: n_threads_ = n_threads; model_loader_.set_n_threads(n_threads); } + void set_segmented_compute_disabled(bool disabled) { + segmented_compute_disabled_ = disabled; + } + void set_prefetch_disabled(bool disabled) { prefetch_disabled_ = disabled; } void set_enable_mmap(bool enable_mmap) { enable_mmap_ = enable_mmap; } void set_writable_mmap(bool writable_mmap) { writable_mmap_ = writable_mmap; } void set_common_ignore_tensors(std::set ignore_tensors); void set_loras(std::vector loras, SDVersion version); - void set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft); + void set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft, const std::vector>& device_limits); static bool tensor_shape_supports_split_buffer(const ggml_tensor* tensor); @@ -199,10 +233,27 @@ public: bool assign_compute_backend(const std::vector& tensors, ggml_backend_t compute_backend) override; bool prepare_params(const std::vector& tensors) override; + void set_workspace_reclaimer(uintptr_t owner_id, std::function reclaim) override; + void remove_runtime_owner(uintptr_t owner_id) override; + bool fits_compute_backend_capacity(const DeviceMemoryRequest& request, + const std::vector& required_params) const override; + bool segmented_compute_enabled() const override { return !segmented_compute_disabled_; } + bool prefetch_enabled() const override { return !prefetch_disabled_; } void release_compute_backend_params(const std::vector& tensors) override; - void release_params_backend_params(const std::vector& tensors) override; - bool prefetch_params(uintptr_t owner_id, - const std::vector& tensors) override; + void evict_compute_backend_params(const std::vector& tensors) override; + WeightResidencyInfo inspect_compute_backend_params( + const std::vector& tensors) const override; + void update_runtime_residency(uintptr_t owner_id, + ggml_backend_t compute_backend, + size_t resident_bytes) override; + bool ensure_compute_backend_capacity( + const DeviceMemoryRequest& request, + const std::vector& required_params, + const std::vector>& preferred_eviction_order, + const std::vector& protected_params) override; + WeightPrefetchResult prefetch_params( + uintptr_t owner_id, + const std::vector& tensors) override; bool activate_prefetched_params(uintptr_t owner_id, const std::vector& tensors) override; void clear_prefetched_params(uintptr_t owner_id) override; diff --git a/src/model_manager_prefetch.cpp b/src/model_manager_prefetch.cpp index 3d39a998..1d304193 100644 --- a/src/model_manager_prefetch.cpp +++ b/src/model_manager_prefetch.cpp @@ -41,15 +41,21 @@ void ModelManager::synchronize_prefetch_block(PrefetchBlock& block) { void ModelManager::free_prefetch_block(PrefetchBlock& block) { synchronize_prefetch_block(block); - block.staged_tensors.clear(); - if (block.buffer != nullptr) { - ggml_backend_buffer_free(block.buffer); - block.buffer = nullptr; - } - if (block.staging_ctx != nullptr) { - ggml_free(block.staging_ctx); - block.staging_ctx = nullptr; + for (auto& staging_block : block.staging_blocks) { + if (staging_block == nullptr) { + continue; + } + staging_block->staged_tensors.clear(); + if (staging_block->buffer != nullptr) { + ggml_backend_buffer_free(staging_block->buffer); + staging_block->buffer = nullptr; + } + if (staging_block->staging_ctx != nullptr) { + ggml_free(staging_block->staging_ctx); + staging_block->staging_ctx = nullptr; + } } + block.staging_blocks.clear(); } bool ModelManager::populate_prefetch_block(PrefetchBlock& block) { @@ -62,26 +68,13 @@ bool ModelManager::populate_prefetch_block(PrefetchBlock& block) { return false; } - ggml_init_params init_params; - init_params.mem_size = block.states.size() * ggml_tensor_overhead(); - init_params.mem_buffer = nullptr; - init_params.no_alloc = true; - block.staging_ctx = ggml_init(init_params); - if (block.staging_ctx == nullptr) { - return false; - } - - block.staged_tensors.reserve(block.states.size()); for (TensorState* state : block.states) { if (state == nullptr || state->tensor == nullptr || state->tensor->buffer == nullptr || state->tensor->data == nullptr || state->params_backend == nullptr || state->staged_to_compute_backend || - state->active_prepare_count > 0) { + state->pin_count > 0) { return false; } - ggml_tensor* staging_tensor = ggml_dup_tensor(block.staging_ctx, state->tensor); - ggml_set_name(staging_tensor, state->tensor->name); - block.staged_tensors.push_back({state, staging_tensor}); } ggml_backend_buffer_type_t buffer_type = @@ -89,38 +82,91 @@ bool ModelManager::populate_prefetch_block(PrefetchBlock& block) { if (buffer_type == nullptr) { return false; } - block.buffer = ggml_backend_alloc_ctx_tensors_from_buft(block.staging_ctx, buffer_type); - if (block.buffer == nullptr) { - return false; + const size_t alignment = ggml_backend_buft_get_alignment(buffer_type); + size_t backend_limit = ggml_backend_buft_get_max_size(buffer_type); + if (!ggml_backend_buft_is_host(buffer_type) && + (backend_limit == 0 || backend_limit > MAX_RESIDENCY_BLOCK_BYTES)) { + backend_limit = MAX_RESIDENCY_BLOCK_BYTES; } - ggml_backend_buffer_set_usage(block.buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - for (const auto& pair : block.staged_tensors) { - TensorState* state = pair.first; - ggml_tensor* staging_tensor = pair.second; - const bool host_source = ggml_backend_buffer_is_host(state->tensor->buffer); - if (host_source && - (!ggml_is_contiguous(state->tensor) || !ggml_is_contiguous(staging_tensor) || - ggml_nbytes(state->tensor) != ggml_nbytes(staging_tensor))) { + auto enqueue_chunk = [&](const std::vector& chunk) -> bool { + if (chunk.empty()) { + return true; + } + ggml_init_params init_params; + init_params.mem_size = chunk.size() * ggml_tensor_overhead(); + init_params.mem_buffer = nullptr; + init_params.no_alloc = true; + ggml_context* staging_ctx = ggml_init(init_params); + if (staging_ctx == nullptr) { return false; } - } - for (const auto& pair : block.staged_tensors) { - TensorState* state = pair.first; - ggml_tensor* staging_tensor = pair.second; - if (ggml_backend_buffer_is_host(state->tensor->buffer)) { - ggml_backend_tensor_set_async(block.transfer_backend, - staging_tensor, - state->tensor->data, - 0, - ggml_nbytes(state->tensor)); - } else { - ggml_backend_tensor_copy_async(state->params_backend, - block.transfer_backend, - state->tensor, - staging_tensor); + auto staging_block = std::make_unique(); + staging_block->compute_backend = block.compute_backend; + staging_block->staging_ctx = staging_ctx; + staging_block->staged_tensors.reserve(chunk.size()); + for (TensorState* state : chunk) { + ggml_tensor* staging_tensor = ggml_dup_tensor(staging_ctx, state->tensor); + ggml_set_name(staging_tensor, state->tensor->name); + if (ggml_backend_buffer_is_host(state->tensor->buffer) && + (!ggml_is_contiguous(state->tensor) || !ggml_is_contiguous(staging_tensor) || + ggml_nbytes(state->tensor) != ggml_nbytes(staging_tensor))) { + ggml_free(staging_ctx); + staging_block->staging_ctx = nullptr; + return false; + } + staging_block->staged_tensors.push_back({state, staging_tensor}); } + staging_block->buffer = + ggml_backend_alloc_ctx_tensors_from_buft(staging_ctx, buffer_type); + if (staging_block->buffer == nullptr) { + ggml_free(staging_ctx); + staging_block->staging_ctx = nullptr; + return false; + } + ggml_backend_buffer_set_usage(staging_block->buffer, + GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + + for (const auto& pair : staging_block->staged_tensors) { + TensorState* state = pair.first; + ggml_tensor* staging_tensor = pair.second; + const bool host_source = ggml_backend_buffer_is_host(state->tensor->buffer); + if (host_source) { + ggml_backend_tensor_set_async(block.transfer_backend, + staging_tensor, + state->tensor->data, + 0, + ggml_nbytes(state->tensor)); + } else { + ggml_backend_tensor_copy_async(state->params_backend, + block.transfer_backend, + state->tensor, + staging_tensor); + } + } + block.staging_blocks.push_back(std::move(staging_block)); + return true; + }; + + std::vector chunk; + size_t chunk_size = 0; + for (TensorState* state : block.states) { + const size_t tensor_size = GGML_PAD( + ggml_backend_buft_get_alloc_size(buffer_type, state->tensor), alignment); + if (!chunk.empty() && backend_limit > 0 && + tensor_size > backend_limit - std::min(chunk_size, backend_limit)) { + if (!enqueue_chunk(chunk)) { + return false; + } + chunk.clear(); + chunk_size = 0; + } + chunk.push_back(state); + chunk_size = tensor_size > SIZE_MAX - chunk_size ? SIZE_MAX : chunk_size + tensor_size; + } + if (!enqueue_chunk(chunk)) { + return false; } ggml_backend_dev_t device = ggml_backend_get_device(block.transfer_backend); @@ -129,48 +175,79 @@ bool ModelManager::populate_prefetch_block(PrefetchBlock& block) { ggml_backend_event_record(block.event, block.transfer_backend); } - LOG_DEBUG("model manager queued layer prefetch (%6.2f MB, %zu tensors) to %s", - ggml_backend_buffer_get_size(block.buffer) / (1024.f * 1024.f), + size_t total_size = 0; + for (const auto& staging_block : block.staging_blocks) { + if (staging_block != nullptr && staging_block->buffer != nullptr) { + const size_t buffer_size = ggml_backend_buffer_get_size(staging_block->buffer); + total_size = buffer_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + buffer_size; + } + } + LOG_DEBUG("model manager queued segment prefetch (%6.2f MB, %zu tensors) to %s", + total_size / (1024.f * 1024.f), block.states.size(), ggml_backend_name(block.compute_backend)); return true; } -bool ModelManager::prefetch_params(uintptr_t owner_id, - const std::vector& tensors) { - clear_prefetched_params(owner_id); +WeightPrefetchResult ModelManager::prefetch_params( + uintptr_t owner_id, + const std::vector& tensors) { if (tensors.empty()) { - return true; + return WeightPrefetchResult::AlreadyResident; } std::vector required_states; - if (!resolve_required_tensor_states(tensors, required_states) || - !load_tensors_to_params_backend(required_states)) { - return false; + if (!resolve_required_tensor_states(tensors, required_states)) { + return WeightPrefetchResult::Failed; } std::vector states; states.reserve(required_states.size()); ggml_backend_t compute_backend = nullptr; + bool needs_synchronous_load = false; for (TensorState* state : required_states) { if (state == nullptr || should_ignore(*state) || - is_optional_missing_tensor(state->name) || - state->compute_backend == state->params_backend || - state->staged_to_compute_backend || state->active_prepare_count > 0) { + is_optional_missing_tensor(state->name)) { continue; } + if (state->compute_backend == state->params_backend) { + needs_synchronous_load = needs_synchronous_load || + !state->loaded_to_params_backend; + continue; + } + if (state->staged_to_compute_backend || state->pin_count > 0) { + continue; + } + // Split buffers cannot use the primary device's asynchronous upload path. + if (split_buffer_type_for(*state) != nullptr) { + return WeightPrefetchResult::Unsupported; + } if (compute_backend == nullptr) { compute_backend = state->compute_backend; } else if (compute_backend != state->compute_backend) { - return false; + return WeightPrefetchResult::Failed; } states.push_back(state); } if (states.empty()) { - return true; + return needs_synchronous_load ? WeightPrefetchResult::Unsupported + : WeightPrefetchResult::AlreadyResident; } if (compute_backend == nullptr || sd_backend_is_cpu(compute_backend)) { - return false; + return WeightPrefetchResult::Unsupported; + } + ggml_backend_dev_t compute_device = ggml_backend_get_device(compute_backend); + ggml_backend_dev_props compute_props{}; + if (compute_device == nullptr) { + return WeightPrefetchResult::Unsupported; + } + ggml_backend_dev_get_props(compute_device, &compute_props); + if (!compute_props.caps.async) { + return WeightPrefetchResult::Unsupported; + } + clear_prefetched_params(owner_id); + if (!load_tensors_to_params_backend(states)) { + return WeightPrefetchResult::Failed; } auto block = std::make_unique(); @@ -178,10 +255,10 @@ bool ModelManager::prefetch_params(uintptr_t owner_id, block->compute_backend = compute_backend; if (!populate_prefetch_block(*block)) { free_prefetch_block(*block); - return false; + return WeightPrefetchResult::Failed; } prefetch_blocks_[owner_id] = std::move(block); - return true; + return WeightPrefetchResult::Scheduled; } bool ModelManager::activate_prefetched_params( @@ -214,32 +291,37 @@ bool ModelManager::activate_prefetched_params( prefetch_blocks_.erase(existing); synchronize_prefetch_block(*block); - for (const auto& pair : block->staged_tensors) { - TensorState* state = pair.first; - ggml_tensor* staging_tensor = pair.second; - if (state == nullptr || state->tensor == nullptr || staging_tensor == nullptr || - state->staged_to_compute_backend || state->active_prepare_count > 0) { - free_prefetch_block(*block); - return false; + for (const auto& staging_block : block->staging_blocks) { + if (staging_block == nullptr) { + continue; + } + for (const auto& pair : staging_block->staged_tensors) { + TensorState* state = pair.first; + ggml_tensor* staging_tensor = pair.second; + if (state == nullptr || state->tensor == nullptr || staging_tensor == nullptr || + state->staged_to_compute_backend || state->pin_count > 0) { + free_prefetch_block(*block); + return false; + } } } - for (auto& pair : block->staged_tensors) { - TensorState* state = pair.first; - ggml_tensor* staging_tensor = pair.second; - std::swap(state->tensor->buffer, staging_tensor->buffer); - std::swap(state->tensor->data, staging_tensor->data); - std::swap(state->tensor->extra, staging_tensor->extra); - state->staged_to_compute_backend = true; + const uint64_t use_epoch = ++residency_epoch_; + for (auto& staging_block : block->staging_blocks) { + if (staging_block == nullptr) { + continue; + } + for (auto& pair : staging_block->staged_tensors) { + TensorState* state = pair.first; + ggml_tensor* staging_tensor = pair.second; + std::swap(state->tensor->buffer, staging_tensor->buffer); + std::swap(state->tensor->data, staging_tensor->data); + std::swap(state->tensor->extra, staging_tensor->extra); + state->staged_to_compute_backend = true; + state->last_use_epoch = use_epoch; + } + compute_staging_blocks_.push_back(std::move(staging_block)); } - - auto staging_block = std::make_unique(); - staging_block->compute_backend = block->compute_backend; - staging_block->buffer = block->buffer; - staging_block->staging_ctx = block->staging_ctx; - staging_block->staged_tensors = std::move(block->staged_tensors); - block->buffer = nullptr; - block->staging_ctx = nullptr; - compute_staging_blocks_.push_back(std::move(staging_block)); + block->staging_blocks.clear(); return true; } diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 217212cc..c0ba9ebd 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -248,9 +248,9 @@ public: sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr}; bool enable_mmap = false; sd::ggml_graph_cut::MaxVramAssignment max_vram_assignment; - bool stream_layers = false; - bool disable_prefetch = false; - bool eager_load = false; + bool disable_prefetch = false; + bool disable_segmented_compute = false; + bool eager_load = false; std::string backend_spec; std::string params_backend_spec; std::string split_mode_spec; @@ -435,7 +435,11 @@ public: if (split_buft == nullptr) { return fall_back_to_layer_split("backend has no split buffer type"); } - model_manager->set_split_buffer_type(main_backend, split_buft); + std::vector> split_device_limits; + for (auto backend : module_backends) { + split_device_limits.emplace_back(backend, max_vram_assignment.bytes_for_backend(backend)); + } + model_manager->set_split_buffer_type(main_backend, split_buft, split_device_limits); std::map split_tensors; if constexpr (std::is_base_of_v) { @@ -599,6 +603,8 @@ public: version, "", model_manager); + control_net->set_max_graph_vram_bytes( + max_graph_vram_bytes_for_module(SDBackendModule::CONTROL_NET)); if (diffusion_conv_direct) { LOG_INFO("Using Conv2d direct in the control net"); control_net->set_conv2d_direct_enabled(true); @@ -860,15 +866,15 @@ public: } bool init(const sd_ctx_params_t* sd_ctx_params) { - n_threads = sd_ctx_params->n_threads; - enable_mmap = sd_ctx_params->enable_mmap; - stream_layers = sd_ctx_params->stream_layers; - disable_prefetch = sd_ctx_params->disable_prefetch; - eager_load = sd_ctx_params->eager_load; - backend_spec = SAFE_STR(sd_ctx_params->backend); - params_backend_spec = SAFE_STR(sd_ctx_params->params_backend); - split_mode_spec = SAFE_STR(sd_ctx_params->split_mode); - auto_fit_enabled = sd_ctx_params->auto_fit; + n_threads = sd_ctx_params->n_threads; + enable_mmap = sd_ctx_params->enable_mmap; + disable_prefetch = sd_ctx_params->disable_prefetch; + disable_segmented_compute = sd_ctx_params->disable_segmented_compute; + eager_load = sd_ctx_params->eager_load; + backend_spec = SAFE_STR(sd_ctx_params->backend); + params_backend_spec = SAFE_STR(sd_ctx_params->params_backend); + split_mode_spec = SAFE_STR(sd_ctx_params->split_mode); + auto_fit_enabled = sd_ctx_params->auto_fit; max_vram_assignment.reset(0.f); { std::string error; @@ -897,6 +903,8 @@ public: model_manager = std::make_shared(); model_manager->set_n_threads(n_threads); model_manager->set_enable_mmap(enable_mmap); + model_manager->set_segmented_compute_disabled(disable_segmented_compute); + model_manager->set_prefetch_disabled(disable_prefetch); ModelLoader& model_loader = model_manager->loader(); if (!init_model_loader(model_loader, sd_ctx_params, use_tae, use_audio_vae, use_control_net)) { @@ -931,10 +939,6 @@ public: return false; } } - if (stream_layers && !backend_manager.params_backend_is_cpu(SDBackendModule::DIFFUSION)) { - LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring"); - stream_layers = false; - } if (eager_load && graph_cut_layer_split_active()) { LOG_WARN("--eager-load is not supported with graph-cut layer split; weights will be prepared lazily"); eager_load = false; @@ -984,7 +988,8 @@ public: } // Avoid full-model LoRA merge buffers on constrained setups. const bool params_offloaded = params_backend_for(SDBackendModule::DIFFUSION) != backend_for(SDBackendModule::DIFFUSION); - const bool streaming_constrained = stream_layers || params_offloaded; + const bool streaming_constrained = params_offloaded || + backend_manager.params_backend_is_disk(SDBackendModule::DIFFUSION); if (have_quantized_weight || streaming_constrained || row_split_active()) { apply_lora_immediately = false; } else { @@ -1357,8 +1362,6 @@ public: } diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION)); - diffusion_model->set_stream_layers_enabled(stream_layers); - diffusion_model->set_layer_prefetch_enabled(!disable_prefetch); if (!register_runner_params("Diffusion model", diffusion_model, SDBackendModule::DIFFUSION, @@ -1368,8 +1371,6 @@ public: if (high_noise_diffusion_model) { high_noise_diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION)); - high_noise_diffusion_model->set_stream_layers_enabled(stream_layers); - high_noise_diffusion_model->set_layer_prefetch_enabled(!disable_prefetch); if (!register_runner_params("High noise diffusion model", high_noise_diffusion_model, SDBackendModule::DIFFUSION, @@ -1574,6 +1575,8 @@ public: version, "", model_manager); + control_net->set_max_graph_vram_bytes( + max_graph_vram_bytes_for_module(SDBackendModule::CONTROL_NET)); if (sd_ctx_params->diffusion_conv_direct) { LOG_INFO("Using Conv2d direct in the control net"); control_net->set_conv2d_direct_enabled(true); @@ -1902,15 +1905,15 @@ public: } bool is_using_v_parameterization_for_sd2(bool is_inpaint = false) { - struct RunnerDoneOnExit { + struct RunnerEndOnExit { GGMLRunner* runner = nullptr; - ~RunnerDoneOnExit() { + ~RunnerEndOnExit() { if (runner != nullptr) { - runner->runner_done(); + runner->runner_end(); } } }; - RunnerDoneOnExit diffusion_runner_done{diffusion_model.get()}; + RunnerEndOnExit diffusion_runner_end{diffusion_model.get()}; sd::Tensor x_t = sd::full({8, 8, 4, 1}, 0.5f); sd::Tensor c = sd::full({1024, 2, 1, 1}, 0.5f); @@ -2541,17 +2544,17 @@ public: const sd_cache_params_t* cache_params, bool preview_final_step, const sd::Tensor& video_positions = {}) { - struct RunnerDoneOnExit { + struct RunnerEndOnExit { GGMLRunner* runner = nullptr; - ~RunnerDoneOnExit() { + ~RunnerEndOnExit() { if (runner != nullptr) { - runner->runner_done(); + runner->runner_end(); } } }; - RunnerDoneOnExit sample_diffusion_runner_done{work_diffusion_model.get()}; + RunnerEndOnExit sample_diffusion_runner_end{work_diffusion_model.get()}; - RunnerDoneOnExit sample_control_runner_done{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr}; + RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr}; std::vector skip_layers(guidance.slg.layers, guidance.slg.layers + guidance.slg.layer_count); float cfg_scale = guidance.txt_cfg; @@ -2923,10 +2926,6 @@ public: LOG_ERROR("Diffusion model sampling failed"); if (control_net) { control_net->free_control_ctx(); - control_net->free_compute_buffer(); - } - if (work_diffusion_model) { - work_diffusion_model->free_compute_buffer(); } return {}; } @@ -2940,10 +2939,6 @@ public: if (control_net) { control_net->free_control_ctx(); - control_net->free_compute_buffer(); - } - if (work_diffusion_model) { - work_diffusion_model->free_compute_buffer(); } return x0; } @@ -3092,7 +3087,6 @@ public: while (decoded.empty() && auto_fit_enabled && sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) { - first_stage_model->free_compute_buffer(); decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); } return decoded; @@ -3567,27 +3561,27 @@ void sd_hires_params_init(sd_hires_params_t* hires_params) { } void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { - *sd_ctx_params = {}; - sd_ctx_params->n_threads = sd_get_num_physical_cores(); - sd_ctx_params->wtype = SD_TYPE_COUNT; - sd_ctx_params->rng_type = CUDA_RNG; - sd_ctx_params->sampler_rng_type = RNG_TYPE_COUNT; - sd_ctx_params->prediction = PREDICTION_COUNT; - sd_ctx_params->lora_apply_mode = LORA_APPLY_AUTO; - sd_ctx_params->max_vram = nullptr; - sd_ctx_params->stream_layers = false; - sd_ctx_params->disable_prefetch = false; - sd_ctx_params->eager_load = false; - sd_ctx_params->enable_mmap = false; - sd_ctx_params->diffusion_flash_attn = false; - sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO; - sd_ctx_params->backend = nullptr; - sd_ctx_params->params_backend = nullptr; - sd_ctx_params->split_mode = nullptr; - sd_ctx_params->auto_fit = false; - sd_ctx_params->rpc_servers = nullptr; - sd_ctx_params->model_args = nullptr; - sd_ctx_params->pulid_weights_path = nullptr; + *sd_ctx_params = {}; + sd_ctx_params->n_threads = sd_get_num_physical_cores(); + sd_ctx_params->wtype = SD_TYPE_COUNT; + sd_ctx_params->rng_type = CUDA_RNG; + sd_ctx_params->sampler_rng_type = RNG_TYPE_COUNT; + sd_ctx_params->prediction = PREDICTION_COUNT; + sd_ctx_params->lora_apply_mode = LORA_APPLY_AUTO; + sd_ctx_params->max_vram = nullptr; + sd_ctx_params->disable_prefetch = false; + sd_ctx_params->disable_segmented_compute = false; + sd_ctx_params->eager_load = false; + sd_ctx_params->enable_mmap = false; + sd_ctx_params->diffusion_flash_attn = false; + sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO; + sd_ctx_params->backend = nullptr; + sd_ctx_params->params_backend = nullptr; + sd_ctx_params->split_mode = nullptr; + sd_ctx_params->auto_fit = false; + sd_ctx_params->rpc_servers = nullptr; + sd_ctx_params->model_args = nullptr; + sd_ctx_params->pulid_weights_path = nullptr; } char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { @@ -3621,8 +3615,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { "sampler_rng_type: %s\n" "prediction: %s\n" "max_vram: %s\n" - "stream_layers: %s\n" "disable_prefetch: %s\n" + "disable_segmented_compute: %s\n" "eager_load: %s\n" "backend: %s\n" "params_backend: %s\n" @@ -3656,8 +3650,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { sd_rng_type_name(sd_ctx_params->sampler_rng_type), sd_prediction_name(sd_ctx_params->prediction), SAFE_STR(sd_ctx_params->max_vram), - BOOL_STR(sd_ctx_params->stream_layers), BOOL_STR(sd_ctx_params->disable_prefetch), + BOOL_STR(sd_ctx_params->disable_segmented_compute), BOOL_STR(sd_ctx_params->eager_load), SAFE_STR(sd_ctx_params->backend), SAFE_STR(sd_ctx_params->params_backend), @@ -4809,11 +4803,11 @@ struct ImageGenerationEmbeds { SDCondition img_uncond; }; -struct ConditionerRunnerDoneOnExit { +struct ConditionerRunnerEndOnExit { Conditioner* conditioner = nullptr; - ~ConditionerRunnerDoneOnExit() { + ~ConditionerRunnerEndOnExit() { if (conditioner != nullptr) { - conditioner->runner_done(); + conditioner->runner_end(); } } }; @@ -5239,7 +5233,7 @@ static std::optional prepare_image_generation_embeds(sd_c SamplePlan* plan, ImageGenerationLatents* latents, const RefImageParams& ref_image_params) { - ConditionerRunnerDoneOnExit conditioner_runner_done{sd_ctx->sd->cond_stage_model.get()}; + ConditionerRunnerEndOnExit conditioner_runner_end{sd_ctx->sd->cond_stage_model.get()}; ConditionerParams condition_params; condition_params.text = request->prompt; @@ -6536,7 +6530,7 @@ static ImageGenerationEmbeds prepare_video_generation_embeds(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* sd_vid_gen_params, const GenerationRequest& request, const ImageGenerationLatents& latents) { - ConditionerRunnerDoneOnExit conditioner_runner_done{sd_ctx->sd->cond_stage_model.get()}; + ConditionerRunnerEndOnExit conditioner_runner_end{sd_ctx->sd->cond_stage_model.get()}; ImageGenerationEmbeds embeds; ConditionerParams condition_params; diff --git a/src/upscaler.cpp b/src/upscaler.cpp index dbb99af3..8be7b948 100644 --- a/src/upscaler.cpp +++ b/src/upscaler.cpp @@ -32,13 +32,6 @@ void UpscalerGGML::set_max_graph_vram_bytes(size_t max_vram_bytes) { } } -void UpscalerGGML::set_stream_layers_enabled(bool enabled) { - stream_layers_enabled = enabled; - if (esrgan_upscaler) { - esrgan_upscaler->set_stream_layers_enabled(enabled); - } -} - bool UpscalerGGML::load_from_file(const std::string& esrgan_path, int n_threads) { ggml_log_set(ggml_log_callback_default, nullptr); @@ -94,7 +87,6 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path, return false; } esrgan_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes); - esrgan_upscaler->set_stream_layers_enabled(stream_layers_enabled); if (direct) { esrgan_upscaler->set_conv2d_direct_enabled(true); } @@ -139,7 +131,7 @@ sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_te false, on_processing); } - esrgan_upscaler->free_compute_buffer(); + esrgan_upscaler->runner_end(); if (upscaled.empty()) { LOG_ERROR("esrgan compute failed"); return {}; diff --git a/src/upscaler.h b/src/upscaler.h index 38150f59..867f6444 100644 --- a/src/upscaler.h +++ b/src/upscaler.h @@ -20,7 +20,6 @@ struct UpscalerGGML { bool direct = false; int tile_size = 128; size_t max_graph_vram_bytes = 0; - bool stream_layers_enabled = false; std::string backend_spec; std::string params_backend_spec; @@ -34,7 +33,6 @@ struct UpscalerGGML { bool load_from_file(const std::string& esrgan_path, int n_threads); void set_max_graph_vram_bytes(size_t max_vram_bytes); - void set_stream_layers_enabled(bool enabled); sd::Tensor upscale_tensor(const sd::Tensor& input_tensor); sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor); }; diff --git a/src/weight_manager.h b/src/weight_manager.h deleted file mode 100644 index c54e33c4..00000000 --- a/src/weight_manager.h +++ /dev/null @@ -1,25 +0,0 @@ -#ifndef __WEIGHT_MANAGER_H__ -#define __WEIGHT_MANAGER_H__ - -#include -#include - -#include "ggml-backend.h" - -struct ggml_tensor; - -struct RunnerWeightManager { - virtual ~RunnerWeightManager() = default; - virtual bool assign_compute_backend(const std::vector& tensors, - ggml_backend_t compute_backend) = 0; - virtual bool prepare_params(const std::vector& tensors) = 0; - virtual void release_compute_backend_params(const std::vector& tensors) = 0; - virtual void release_params_backend_params(const std::vector& tensors) = 0; - virtual bool prefetch_params(uintptr_t owner_id, - const std::vector& tensors) = 0; - virtual bool activate_prefetched_params(uintptr_t owner_id, - const std::vector& tensors) = 0; - virtual void clear_prefetched_params(uintptr_t owner_id) = 0; -}; - -#endif // __WEIGHT_MANAGER_H__