diff --git a/examples/cli/README.md b/examples/cli/README.md index 7b620fee7..b32fe37f9 100644 --- a/examples/cli/README.md +++ b/examples/cli/README.md @@ -54,6 +54,8 @@ Context Options: -t, --threads number of threads to use during computation (default: -1). If threads <= 0, then threads will be set to the number of CPU physical cores --chroma-t5-mask-pad t5 mask pad size of chroma + --max-vram maximum VRAM budget in GiB for graph-cut segmented execution. 0 disables + graph splitting --force-sdxl-vae-conv-scale force use of conv scale on sdxl vae --offload-to-cpu place the weights in RAM to save VRAM, and automatically load them into VRAM when needed diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 1a5399b82..d4c8a72b8 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -394,7 +394,12 @@ ArgOptions SDContextParams::get_options() { &chroma_t5_mask_pad}, }; - options.float_options = {}; + options.float_options = { + {"", + "--max-vram", + "maximum VRAM budget in GiB for graph-cut segmented execution. 0 disables graph splitting", + &max_vram}, + }; options.bool_options = { {"", @@ -670,6 +675,7 @@ std::string SDContextParams::to_string() const { << " rng_type: " << sd_rng_type_name(rng_type) << ",\n" << " 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" << " enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n" << " control_net_cpu: " << (control_net_cpu ? "true" : "false") << ",\n" << " clip_on_cpu: " << (clip_on_cpu ? "true" : "false") << ",\n" @@ -744,6 +750,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool vae_decode_only, bool f chroma_use_t5_mask, chroma_t5_mask_pad, qwen_image_zero_cond_t, + max_vram, }; return sd_ctx_params; } diff --git a/examples/common/common.h b/examples/common/common.h index c4498c352..f87293f3e 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -109,6 +109,7 @@ struct SDContextParams { rng_type_t rng_type = CUDA_RNG; rng_type_t sampler_rng_type = RNG_TYPE_COUNT; bool offload_params_to_cpu = false; + float max_vram = 0.f; bool enable_mmap = false; bool control_net_cpu = false; bool clip_on_cpu = false; diff --git a/examples/server/README.md b/examples/server/README.md index 469dd346f..23b79c9d8 100644 --- a/examples/server/README.md +++ b/examples/server/README.md @@ -156,6 +156,8 @@ Context Options: -t, --threads number of threads to use during computation (default: -1). If threads <= 0, then threads will be set to the number of CPU physical cores --chroma-t5-mask-pad t5 mask pad size of chroma + --max-vram maximum VRAM budget in GiB for graph-cut segmented execution. 0 disables + graph splitting --force-sdxl-vae-conv-scale force use of conv scale on sdxl vae --offload-to-cpu place the weights in RAM to save VRAM, and automatically load them into VRAM when needed diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 75027f8f8..c4c14949c 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -203,6 +203,7 @@ typedef struct { bool chroma_use_t5_mask; int chroma_t5_mask_pad; bool qwen_image_zero_cond_t; + float max_vram; } sd_ctx_params_t; typedef struct { diff --git a/src/anima.hpp b/src/anima.hpp index 5850cc3e6..4bfc04749 100644 --- a/src/anima.hpp +++ b/src/anima.hpp @@ -499,9 +499,15 @@ namespace Anima { encoder_hidden_states = adapted_context; } + sd::ggml_graph_cut::mark_graph_cut(x, "anima.prelude", "x"); + sd::ggml_graph_cut::mark_graph_cut(embedded_timestep, "anima.prelude", "embedded_timestep"); + sd::ggml_graph_cut::mark_graph_cut(temb, "anima.prelude", "temb"); + sd::ggml_graph_cut::mark_graph_cut(encoder_hidden_states, "anima.prelude", "context"); + for (int i = 0; i < num_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); x = block->forward(ctx, x, encoder_hidden_states, embedded_timestep, temb, image_pe); + sd::ggml_graph_cut::mark_graph_cut(x, "anima.blocks." + std::to_string(i), "x"); } x = final_layer->forward(ctx, x, embedded_timestep, temb); // [N, h*w, ph*pw*C] diff --git a/src/auto_encoder_kl.hpp b/src/auto_encoder_kl.hpp index 5cf09b883..4fb28a16f 100644 --- a/src/auto_encoder_kl.hpp +++ b/src/auto_encoder_kl.hpp @@ -328,6 +328,7 @@ class Encoder : public GGMLBlock { auto conv_out = std::dynamic_pointer_cast(blocks["conv_out"]); auto h = conv_in->forward(ctx, x); // [N, ch, h, w] + // sd::ggml_graph_cut::mark_graph_cut(h, "vae.encoder.prelude", "h"); // downsampling size_t num_resolutions = ch_mult.size(); @@ -337,12 +338,14 @@ class Encoder : public GGMLBlock { auto down_block = std::dynamic_pointer_cast(blocks[name]); h = down_block->forward(ctx, h); + // sd::ggml_graph_cut::mark_graph_cut(h, "vae.encoder.down." + std::to_string(i) + ".block." + std::to_string(j), "h"); } if (i != num_resolutions - 1) { std::string name = "down." + std::to_string(i) + ".downsample"; auto down_sample = std::dynamic_pointer_cast(blocks[name]); h = down_sample->forward(ctx, h); + // sd::ggml_graph_cut::mark_graph_cut(h, "vae.encoder.down." + std::to_string(i) + ".downsample", "h"); } } @@ -350,6 +353,7 @@ class Encoder : public GGMLBlock { h = mid_block_1->forward(ctx, h); h = mid_attn_1->forward(ctx, h); h = mid_block_2->forward(ctx, h); // [N, block_in, h, w] + // sd::ggml_graph_cut::mark_graph_cut(h, "vae.encoder.mid", "h"); // end h = norm_out->forward(ctx, h); @@ -450,6 +454,7 @@ class Decoder : public GGMLBlock { // conv_in auto h = conv_in->forward(ctx, z); // [N, block_in, h, w] + // sd::ggml_graph_cut::mark_graph_cut(h, "vae.decoder.prelude", "h"); // middle h = mid_block_1->forward(ctx, h); @@ -457,6 +462,7 @@ class Decoder : public GGMLBlock { h = mid_attn_1->forward(ctx, h); h = mid_block_2->forward(ctx, h); // [N, block_in, h, w] + // sd::ggml_graph_cut::mark_graph_cut(h, "vae.decoder.mid", "h"); // upsampling int num_resolutions = static_cast(ch_mult.size()); @@ -466,12 +472,14 @@ class Decoder : public GGMLBlock { auto up_block = std::dynamic_pointer_cast(blocks[name]); h = up_block->forward(ctx, h); + // sd::ggml_graph_cut::mark_graph_cut(h, "vae.decoder.up." + std::to_string(i) + ".block." + std::to_string(j), "h"); } if (i != 0) { std::string name = "up." + std::to_string(i) + ".upsample"; auto up_sample = std::dynamic_pointer_cast(blocks[name]); h = up_sample->forward(ctx, h); + // sd::ggml_graph_cut::mark_graph_cut(h, "vae.decoder.up." + std::to_string(i) + ".upsample", "h"); } } @@ -599,6 +607,7 @@ class AutoEncoderKLModel : public GGMLBlock { if (use_quant) { auto post_quant_conv = std::dynamic_pointer_cast(blocks["post_quant_conv"]); z = post_quant_conv->forward(ctx, z); // [N, z_channels, h, w] + // sd::ggml_graph_cut::mark_graph_cut(z, "vae.decode.prelude", "z"); } auto decoder = std::dynamic_pointer_cast(blocks["decoder"]); @@ -616,6 +625,7 @@ class AutoEncoderKLModel : public GGMLBlock { if (use_quant) { auto quant_conv = std::dynamic_pointer_cast(blocks["quant_conv"]); z = quant_conv->forward(ctx, z); // [N, 2*embed_dim, h/8, w/8] + // sd::ggml_graph_cut::mark_graph_cut(z, "vae.encode.final", "z"); } if (sd_version_uses_flux2_vae(version)) { z = ggml_ext_chunk(ctx->ggml_ctx, z, 2, 2)[0]; diff --git a/src/clip.hpp b/src/clip.hpp index 8a2070e0b..8b2084c49 100644 --- a/src/clip.hpp +++ b/src/clip.hpp @@ -95,8 +95,9 @@ struct CLIPEncoder : public GGMLBlock { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, - ggml_tensor* mask = nullptr, - int clip_skip = -1) { + ggml_tensor* mask = nullptr, + int clip_skip = -1, + const std::string& graph_cut_prefix = "") { // x: [N, n_token, d_model] int layer_idx = n_layer - 1; // LOG_DEBUG("clip_skip %d", clip_skip); @@ -112,6 +113,9 @@ struct CLIPEncoder : public GGMLBlock { std::string name = "layers." + std::to_string(i); auto layer = std::dynamic_pointer_cast(blocks[name]); x = layer->forward(ctx, x, mask); // [N, n_token, d_model] + if (!graph_cut_prefix.empty()) { + sd::ggml_graph_cut::mark_graph_cut(x, graph_cut_prefix + ".layers." + std::to_string(i), "x"); + } // LOG_DEBUG("layer %d", i); } return x; @@ -304,7 +308,8 @@ class CLIPTextModel : public GGMLBlock { auto final_layer_norm = std::dynamic_pointer_cast(blocks["final_layer_norm"]); auto x = embeddings->forward(ctx, input_ids, tkn_embeddings); // [N, n_token, hidden_size] - x = encoder->forward(ctx, x, mask, return_pooled ? -1 : clip_skip); + sd::ggml_graph_cut::mark_graph_cut(x, "clip_text.prelude", "x"); + x = encoder->forward(ctx, x, mask, return_pooled ? -1 : clip_skip, "clip_text"); if (return_pooled || with_final_ln) { x = final_layer_norm->forward(ctx, x); } @@ -368,7 +373,8 @@ class CLIPVisionModel : public GGMLBlock { auto x = embeddings->forward(ctx, pixel_values); // [N, num_positions, embed_dim] x = pre_layernorm->forward(ctx, x); - x = encoder->forward(ctx, x, nullptr, clip_skip); + sd::ggml_graph_cut::mark_graph_cut(x, "clip_vision.prelude", "x"); + x = encoder->forward(ctx, x, nullptr, clip_skip, "clip_vision"); auto last_hidden_state = x; diff --git a/src/conditioner.hpp b/src/conditioner.hpp index 9f4d45524..4907938b0 100644 --- a/src/conditioner.hpp +++ b/src/conditioner.hpp @@ -85,7 +85,8 @@ struct Conditioner { virtual void free_params_buffer() = 0; virtual void get_param_tensors(std::map& tensors) = 0; virtual size_t get_params_buffer_size() = 0; - virtual void set_flash_attention_enabled(bool enabled) = 0; + virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {} + virtual void set_flash_attention_enabled(bool enabled) = 0; virtual void set_weight_adapter(const std::shared_ptr& adapter) {} virtual std::tuple> get_learned_condition_with_trigger(int n_threads, const ConditionerParams& conditioner_params) { @@ -165,6 +166,13 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { return buffer_size; } + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { + text_model->set_max_graph_vram_bytes(max_vram_bytes); + if (sd_version_is_sdxl(version)) { + text_model2->set_max_graph_vram_bytes(max_vram_bytes); + } + } + void set_flash_attention_enabled(bool enabled) override { text_model->set_flash_attention_enabled(enabled); if (sd_version_is_sdxl(version)) { @@ -781,6 +789,18 @@ struct SD3CLIPEmbedder : public Conditioner { return buffer_size; } + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { + if (clip_l) { + clip_l->set_max_graph_vram_bytes(max_vram_bytes); + } + if (clip_g) { + clip_g->set_max_graph_vram_bytes(max_vram_bytes); + } + if (t5) { + t5->set_max_graph_vram_bytes(max_vram_bytes); + } + } + void set_flash_attention_enabled(bool enabled) override { if (clip_l) { clip_l->set_flash_attention_enabled(enabled); @@ -1124,6 +1144,15 @@ struct FluxCLIPEmbedder : public Conditioner { return buffer_size; } + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { + if (clip_l) { + clip_l->set_max_graph_vram_bytes(max_vram_bytes); + } + if (t5) { + t5->set_max_graph_vram_bytes(max_vram_bytes); + } + } + void set_flash_attention_enabled(bool enabled) override { if (clip_l) { clip_l->set_flash_attention_enabled(enabled); @@ -1349,6 +1378,12 @@ struct T5CLIPEmbedder : public Conditioner { return buffer_size; } + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { + if (t5) { + t5->set_max_graph_vram_bytes(max_vram_bytes); + } + } + void set_flash_attention_enabled(bool enabled) override { if (t5) { t5->set_flash_attention_enabled(enabled); @@ -1525,6 +1560,10 @@ struct AnimaConditioner : public Conditioner { return llm->get_params_buffer_size(); } + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { + llm->set_max_graph_vram_bytes(max_vram_bytes); + } + void set_flash_attention_enabled(bool enabled) override { llm->set_flash_attention_enabled(enabled); } @@ -1657,6 +1696,10 @@ struct LLMEmbedder : public Conditioner { return buffer_size; } + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { + llm->set_max_graph_vram_bytes(max_vram_bytes); + } + void set_flash_attention_enabled(bool enabled) override { llm->set_flash_attention_enabled(enabled); } diff --git a/src/diffusion_model.hpp b/src/diffusion_model.hpp index c0a2a11c0..1a202a1a7 100644 --- a/src/diffusion_model.hpp +++ b/src/diffusion_model.hpp @@ -49,6 +49,7 @@ struct DiffusionModel { virtual void set_weight_adapter(const std::shared_ptr& adapter){}; virtual int64_t get_adm_in_channels() = 0; virtual void set_flash_attention_enabled(bool enabled) = 0; + virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) = 0; virtual void set_circular_axes(bool circular_x, bool circular_y) = 0; }; @@ -98,6 +99,10 @@ struct UNetModel : public DiffusionModel { unet.set_flash_attention_enabled(enabled); } + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { + unet.set_max_graph_vram_bytes(max_vram_bytes); + } + void set_circular_axes(bool circular_x, bool circular_y) override { unet.set_circular_axes(circular_x, circular_y); } @@ -164,6 +169,10 @@ struct MMDiTModel : public DiffusionModel { mmdit.set_flash_attention_enabled(enabled); } + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { + mmdit.set_max_graph_vram_bytes(max_vram_bytes); + } + void set_circular_axes(bool circular_x, bool circular_y) override { mmdit.set_circular_axes(circular_x, circular_y); } @@ -229,6 +238,10 @@ struct FluxModel : public DiffusionModel { flux.set_flash_attention_enabled(enabled); } + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { + flux.set_max_graph_vram_bytes(max_vram_bytes); + } + void set_circular_axes(bool circular_x, bool circular_y) override { flux.set_circular_axes(circular_x, circular_y); } @@ -299,6 +312,10 @@ struct AnimaModel : public DiffusionModel { anima.set_flash_attention_enabled(enabled); } + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { + anima.set_max_graph_vram_bytes(max_vram_bytes); + } + void set_circular_axes(bool circular_x, bool circular_y) override { anima.set_circular_axes(circular_x, circular_y); } @@ -364,6 +381,10 @@ struct WanModel : public DiffusionModel { wan.set_flash_attention_enabled(enabled); } + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { + wan.set_max_graph_vram_bytes(max_vram_bytes); + } + void set_circular_axes(bool circular_x, bool circular_y) override { wan.set_circular_axes(circular_x, circular_y); } @@ -433,6 +454,10 @@ struct QwenImageModel : public DiffusionModel { qwen_image.set_flash_attention_enabled(enabled); } + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { + qwen_image.set_max_graph_vram_bytes(max_vram_bytes); + } + void set_circular_axes(bool circular_x, bool circular_y) override { qwen_image.set_circular_axes(circular_x, circular_y); } @@ -499,6 +524,10 @@ struct ZImageModel : public DiffusionModel { z_image.set_flash_attention_enabled(enabled); } + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { + z_image.set_max_graph_vram_bytes(max_vram_bytes); + } + void set_circular_axes(bool circular_x, bool circular_y) override { z_image.set_circular_axes(circular_x, circular_y); } @@ -564,6 +593,10 @@ struct ErnieImageModel : public DiffusionModel { ernie_image.set_flash_attention_enabled(enabled); } + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { + ernie_image.set_max_graph_vram_bytes(max_vram_bytes); + } + void set_circular_axes(bool circular_x, bool circular_y) override { ernie_image.set_circular_axes(circular_x, circular_y); } diff --git a/src/ernie_image.hpp b/src/ernie_image.hpp index d17648d2b..931794f1a 100644 --- a/src/ernie_image.hpp +++ b/src/ernie_image.hpp @@ -295,7 +295,9 @@ namespace ErnieImage { auto c = time_embedding->forward(ctx, sample); // [N, hidden_size] auto mod_params = adaLN_mod->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); // [N, 6 * hidden_size] - auto chunks = ggml_ext_chunk(ctx->ggml_ctx, mod_params, 6, 0); + sd::ggml_graph_cut::mark_graph_cut(hidden_states, "ernie_image.prelude", "hidden_states"); + // sd::ggml_graph_cut::mark_graph_cut(mod_params, "ernie_image.prelude", "mod_params"); + auto chunks = ggml_ext_chunk(ctx->ggml_ctx, mod_params, 6, 0); std::vector temb; temb.reserve(6); for (auto chunk : chunks) { @@ -305,6 +307,7 @@ namespace ErnieImage { for (int i = 0; i < params.num_layers; i++) { auto layer = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); hidden_states = layer->forward(ctx, hidden_states, pe, temb); + sd::ggml_graph_cut::mark_graph_cut(hidden_states, "ernie_image.layers." + std::to_string(i), "hidden_states"); } hidden_states = final_norm->forward(ctx, hidden_states, c); diff --git a/src/esrgan.hpp b/src/esrgan.hpp index 26c46f5b3..f84b77a29 100644 --- a/src/esrgan.hpp +++ b/src/esrgan.hpp @@ -124,27 +124,33 @@ class RRDBNet : public GGMLBlock { auto conv_hr = std::dynamic_pointer_cast(blocks["conv_hr"]); auto conv_last = std::dynamic_pointer_cast(blocks["conv_last"]); - auto feat = conv_first->forward(ctx, x); + auto feat = conv_first->forward(ctx, x); + sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.prelude", "feat"); auto body_feat = feat; for (int i = 0; i < num_block; i++) { std::string name = "body." + std::to_string(i); auto block = std::dynamic_pointer_cast(blocks[name]); body_feat = block->forward(ctx, body_feat); + sd::ggml_graph_cut::mark_graph_cut(body_feat, "esrgan.body." + std::to_string(i), "feat"); } body_feat = conv_body->forward(ctx, body_feat); feat = ggml_add(ctx->ggml_ctx, feat, body_feat); + sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.body.out", "feat"); // upsample if (scale >= 2) { auto conv_up1 = std::dynamic_pointer_cast(blocks["conv_up1"]); feat = lrelu(ctx, conv_up1->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST))); + sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.up1", "feat"); if (scale == 4) { auto conv_up2 = std::dynamic_pointer_cast(blocks["conv_up2"]); feat = lrelu(ctx, conv_up2->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST))); + sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.up2", "feat"); } } // for all scales auto out = conv_last->forward(ctx, lrelu(ctx, conv_hr->forward(ctx, feat))); + sd::ggml_graph_cut::mark_graph_cut(out, "esrgan.final", "out"); return out; } }; diff --git a/src/flux.hpp b/src/flux.hpp index e6bf002fb..732a37197 100644 --- a/src/flux.hpp +++ b/src/flux.hpp @@ -928,6 +928,9 @@ namespace Flux { } txt = txt_in->forward(ctx, txt); + sd::ggml_graph_cut::mark_graph_cut(img, "flux.prelude", "img"); + sd::ggml_graph_cut::mark_graph_cut(txt, "flux.prelude", "txt"); + sd::ggml_graph_cut::mark_graph_cut(vec, "flux.prelude", "vec"); for (int i = 0; i < params.depth; i++) { if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i) != skip_layers.end()) { @@ -939,6 +942,8 @@ namespace Flux { auto img_txt = block->forward(ctx, img, txt, vec, pe, txt_img_mask, ds_img_mods, ds_txt_mods); img = img_txt.first; // [N, n_img_token, hidden_size] txt = img_txt.second; // [N, n_txt_token, hidden_size] + sd::ggml_graph_cut::mark_graph_cut(img, "flux.double_blocks." + std::to_string(i), "img"); + sd::ggml_graph_cut::mark_graph_cut(txt, "flux.double_blocks." + std::to_string(i), "txt"); } auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_img_token, hidden_size] @@ -949,6 +954,7 @@ namespace Flux { auto block = std::dynamic_pointer_cast(blocks["single_blocks." + std::to_string(i)]); txt_img = block->forward(ctx, txt_img, vec, pe, txt_img_mask, ss_mods); + sd::ggml_graph_cut::mark_graph_cut(txt_img, "flux.single_blocks." + std::to_string(i), "txt_img"); } img = ggml_view_3d(ctx->ggml_ctx, diff --git a/src/ggml_extend.hpp b/src/ggml_extend.hpp index 8b748194f..362303229 100644 --- a/src/ggml_extend.hpp +++ b/src/ggml_extend.hpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -26,6 +27,7 @@ #include "ggml-backend.h" #include "ggml.h" #include "ggml_extend_backend.hpp" +#include "ggml_graph_cut.h" #include "model.h" #include "tensor.hpp" @@ -1708,6 +1710,8 @@ struct GGMLRunnerContext { struct GGMLRunner { protected: typedef std::function get_graph_cb_t; + using GraphCutSegment = sd::ggml_graph_cut::Segment; + using GraphCutPlan = sd::ggml_graph_cut::Plan; ggml_backend_t params_backend = nullptr; ggml_backend_t runtime_backend = nullptr; @@ -1724,6 +1728,11 @@ struct GGMLRunner { ggml_context* compute_ctx = nullptr; ggml_gallocr* compute_allocr = nullptr; + ggml_context* partial_offload_ctx = nullptr; + ggml_backend_buffer_t partial_runtime_params_buffer = nullptr; + std::vector> partial_offload_pairs; + size_t max_graph_vram_bytes = 0; + std::shared_ptr weight_adapter = nullptr; std::vector one_vec = {1.f}; @@ -1741,6 +1750,9 @@ struct GGMLRunner { bool circular_x_enabled = false; bool circular_y_enabled = false; + sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_; + std::unordered_set params_tensor_set_; + template static sd::Tensor take_or_empty(std::optional> tensor) { if (!tensor.has_value()) { @@ -1775,6 +1787,7 @@ struct GGMLRunner { params_ctx = ggml_init(params); GGML_ASSERT(params_ctx != nullptr); + params_tensor_set_.clear(); if (params_backend != runtime_backend) { offload_ctx = ggml_init(params); GGML_ASSERT(offload_ctx != nullptr); @@ -1786,10 +1799,15 @@ struct GGMLRunner { ggml_free(params_ctx); params_ctx = nullptr; } + params_tensor_set_.clear(); if (offload_ctx != nullptr) { ggml_free(offload_ctx); offload_ctx = nullptr; } + if (partial_offload_ctx != nullptr) { + ggml_free(partial_offload_ctx); + partial_offload_ctx = nullptr; + } } void alloc_cache_ctx() { @@ -1824,6 +1842,17 @@ struct GGMLRunner { ggml_free(compute_ctx); compute_ctx = nullptr; } + backend_tensor_data_map.clear(); + } + + void rebuild_params_tensor_set() { + params_tensor_set_.clear(); + if (params_ctx == nullptr) { + return; + } + for (ggml_tensor* t = ggml_get_first_tensor(params_ctx); t != nullptr; t = ggml_get_next_tensor(params_ctx, t)) { + params_tensor_set_.insert(t); + } } void prepare_build_in_tensor_before() { @@ -1859,13 +1888,25 @@ struct GGMLRunner { return gf; } - bool alloc_compute_buffer(get_graph_cb_t get_graph) { + bool prepare_compute_graph(get_graph_cb_t get_graph, + ggml_cgraph** gf_out) { + GGML_ASSERT(gf_out != nullptr); + + reset_compute_ctx(); + ggml_cgraph* gf = get_compute_graph(get_graph); + if (gf == nullptr) { + free_compute_ctx(); + return false; + } + + *gf_out = gf; + return true; + } + + bool alloc_compute_buffer(ggml_cgraph* gf) { if (compute_allocr != nullptr) { return true; } - reset_compute_ctx(); - ggml_cgraph* gf = get_compute_graph(get_graph); - backend_tensor_data_map.clear(); compute_allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_backend)); if (!ggml_gallocr_reserve(compute_allocr, gf)) { @@ -1891,47 +1932,132 @@ struct GGMLRunner { } } - void copy_cache_tensors_to_cache_buffer() { - if (cache_tensor_map.size() == 0) { - return; + bool copy_cache_tensors_to_cache_buffer(const std::unordered_set* cache_keep_names = nullptr) { + 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; + } } - free_cache_ctx_and_buffer(); + 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(); - GGML_ASSERT(cache_buffer == nullptr); - std::map runtime_tensor_to_cache_tensor; - for (auto kv : cache_tensor_map) { - auto cache_tensor = ggml_dup_tensor(cache_ctx, kv.second); + 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()); - runtime_tensor_to_cache_tensor[kv.second] = cache_tensor; + 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 (auto kv : runtime_tensor_to_cache_tensor) { - ggml_backend_tensor_copy(kv.first, kv.second); + 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 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 ? 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); - cache_tensor_map.clear(); 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), ggml_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; } - void copy_data_to_backend_tensor() { + void copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy = true) { + GGML_ASSERT(gf != nullptr); + std::unordered_set graph_tensor_set; + const int n_leafs = sd::ggml_graph_cut::leaf_count(gf); + const int n_nodes = ggml_graph_n_nodes(gf); + graph_tensor_set.reserve(static_cast(n_leafs + n_nodes)); + for (int i = 0; i < n_leafs; ++i) { + graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i)); + } + for (int i = 0; i < n_nodes; ++i) { + graph_tensor_set.insert(ggml_graph_node(gf, i)); + } + for (auto& kv : backend_tensor_data_map) { auto tensor = kv.first; auto data = kv.second; + if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) { + continue; + } + + ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + if (buf == nullptr) { + LOG_WARN("%s graph exec skip tensor copy: name=%s op=%s reason=buffer_not_set data=%p view_src=%p view_src_buffer=%p", + get_desc().c_str(), + tensor && tensor->name[0] != '\0' ? tensor->name : "", + tensor ? ggml_op_name(tensor->op) : "", + data, + tensor ? tensor->view_src : nullptr, + (tensor && tensor->view_src) ? tensor->view_src->buffer : nullptr); + continue; + } + ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor)); } - backend_tensor_data_map.clear(); + if (clear_after_copy) { + backend_tensor_data_map.clear(); + } } - bool offload_params_to_runtime_backend() { + bool offload_all_params() { + restore_partial_params(); if (params_backend == runtime_backend) { return true; } @@ -1958,6 +2084,7 @@ struct GGMLRunner { num_tensors); return false; } + ggml_backend_buffer_set_usage(runtime_params_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); ggml_tensor* t = ggml_get_first_tensor(params_ctx); ggml_tensor* offload_t = ggml_get_first_tensor(offload_ctx); @@ -1987,7 +2114,85 @@ struct GGMLRunner { return true; } - void offload_params_to_params_backend() { + bool offload_partial_params(const std::vector& tensors) { + restore_partial_params(); + if (params_backend == runtime_backend) { + return true; + } + if (tensors.empty()) { + return true; + } + GGML_ASSERT(!params_on_runtime_backend); + GGML_ASSERT(partial_runtime_params_buffer == nullptr); + + std::vector unique_tensors; + std::unordered_set seen_tensors; + unique_tensors.reserve(tensors.size()); + seen_tensors.reserve(tensors.size()); + for (ggml_tensor* tensor : tensors) { + if (tensor == nullptr) { + continue; + } + if (seen_tensors.insert(tensor).second) { + unique_tensors.push_back(tensor); + } + } + if (unique_tensors.empty()) { + return true; + } + + ggml_init_params params; + params.mem_size = std::max(1, unique_tensors.size()) * ggml_tensor_overhead(); + params.mem_buffer = nullptr; + params.no_alloc = true; + + partial_offload_ctx = ggml_init(params); + GGML_ASSERT(partial_offload_ctx != nullptr); + + partial_offload_pairs.clear(); + partial_offload_pairs.reserve(unique_tensors.size()); + + for (ggml_tensor* tensor : unique_tensors) { + GGML_ASSERT(tensor->view_src == nullptr); + ggml_tensor* offload_tensor = ggml_dup_tensor(partial_offload_ctx, tensor); + ggml_set_name(offload_tensor, tensor->name); + partial_offload_pairs.push_back({tensor, offload_tensor}); + } + + partial_runtime_params_buffer = ggml_backend_alloc_ctx_tensors(partial_offload_ctx, runtime_backend); + if (partial_runtime_params_buffer == nullptr) { + LOG_ERROR("%s alloc partial runtime params backend buffer failed, num_tensors = %zu", + get_desc().c_str(), + partial_offload_pairs.size()); + ggml_free(partial_offload_ctx); + partial_offload_ctx = nullptr; + partial_offload_pairs.clear(); + return false; + } + ggml_backend_buffer_set_usage(partial_runtime_params_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + + for (auto& pair : partial_offload_pairs) { + ggml_tensor* tensor = pair.first; + ggml_tensor* offload_tensor = pair.second; + + ggml_backend_tensor_copy(tensor, offload_tensor); + std::swap(tensor->buffer, offload_tensor->buffer); + std::swap(tensor->data, offload_tensor->data); + std::swap(tensor->extra, offload_tensor->extra); + } + + size_t params_buffer_size = ggml_backend_buffer_get_size(partial_runtime_params_buffer); + LOG_DEBUG("%s offload partial params (%6.2f MB, %zu tensors) to runtime backend (%s)", + get_desc().c_str(), + params_buffer_size / (1024.f * 1024.f), + partial_offload_pairs.size(), + ggml_backend_name(runtime_backend)); + + return true; + } + + void restore_all_params() { + restore_partial_params(); if (!params_on_runtime_backend) { return; } @@ -2013,17 +2218,323 @@ struct GGMLRunner { params_on_runtime_backend = false; } + void restore_partial_params() { + if (partial_offload_pairs.empty()) { + if (partial_runtime_params_buffer != nullptr) { + ggml_backend_buffer_free(partial_runtime_params_buffer); + partial_runtime_params_buffer = nullptr; + } + if (partial_offload_ctx != nullptr) { + ggml_free(partial_offload_ctx); + partial_offload_ctx = nullptr; + } + return; + } + + for (auto& pair : partial_offload_pairs) { + ggml_tensor* tensor = pair.first; + ggml_tensor* offload_tensor = pair.second; + + tensor->buffer = offload_tensor->buffer; + tensor->data = offload_tensor->data; + tensor->extra = offload_tensor->extra; + offload_tensor->buffer = nullptr; + offload_tensor->data = nullptr; + offload_tensor->extra = nullptr; + } + + if (partial_runtime_params_buffer != nullptr) { + ggml_backend_buffer_free(partial_runtime_params_buffer); + partial_runtime_params_buffer = nullptr; + } + partial_offload_pairs.clear(); + + if (partial_offload_ctx != nullptr) { + ggml_free(partial_offload_ctx); + partial_offload_ctx = nullptr; + } + } + + 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 && + params_backend != runtime_backend && + !ggml_backend_is_cpu(runtime_backend); + } + + bool can_attempt_graph_cut_segmented_compute() const { + return max_graph_vram_bytes > 0 && + params_backend != runtime_backend && + !ggml_backend_is_cpu(runtime_backend); + } + + bool resolve_graph_cut_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_, + max_graph_vram_bytes, + params_tensor_set_, + get_desc().c_str()); + return true; + } + + void reset_segment_runtime_tensors(const GraphCutSegment& segment, + ggml_cgraph* gf) { + 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: + case GraphCutSegment::INPUT_EXTERNAL: + input_tensor->buffer = nullptr; + input_tensor->data = nullptr; + input_tensor->extra = nullptr; + break; + 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; + } + } + return true; + } + + template + std::optional> execute_graph(ggml_cgraph* gf, + int n_threads, + bool free_compute_buffer_immediately, + const std::vector& runtime_param_tensors, + bool preserve_backend_tensor_data_map, + bool no_return = false, + const std::unordered_set* cache_keep_names = nullptr) { + int64_t t_execute_begin = ggml_time_ms(); + const bool use_partial_param_offload = !runtime_param_tensors.empty(); + int64_t t_offload_begin = ggml_time_ms(); + if (use_partial_param_offload) { + if (!offload_partial_params(runtime_param_tensors)) { + LOG_ERROR("%s offload partial params to runtime backend failed", get_desc().c_str()); + return std::nullopt; + } + } else { + if (!offload_all_params()) { + LOG_ERROR("%s offload params to runtime backend failed", get_desc().c_str()); + return std::nullopt; + } + } + int64_t t_offload_end = ggml_time_ms(); + + int64_t t_alloc_begin = ggml_time_ms(); + if (!alloc_compute_buffer(gf)) { + LOG_ERROR("%s alloc compute buffer failed", get_desc().c_str()); + if (use_partial_param_offload) { + restore_partial_params(); + } + return std::nullopt; + } + + if (!ggml_gallocr_alloc_graph(compute_allocr, gf)) { + LOG_ERROR("%s alloc compute graph failed", get_desc().c_str()); + if (free_compute_buffer_immediately) { + free_compute_buffer(); + } else if (use_partial_param_offload) { + restore_partial_params(); + } + return std::nullopt; + } + int64_t t_alloc_end = ggml_time_ms(); + + int64_t t_copy_begin = ggml_time_ms(); + copy_data_to_backend_tensor(gf, !preserve_backend_tensor_data_map); + int64_t t_copy_end = ggml_time_ms(); + if (ggml_backend_is_cpu(runtime_backend)) { + ggml_backend_cpu_set_n_threads(runtime_backend, n_threads); + } + + int64_t t_compute_begin = ggml_time_ms(); + ggml_status status = ggml_backend_graph_compute(runtime_backend, gf); + int64_t t_compute_end = ggml_time_ms(); + if (status != GGML_STATUS_SUCCESS) { + LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status)); + if (free_compute_buffer_immediately) { + free_compute_buffer(); + } else if (use_partial_param_offload) { + restore_partial_params(); + } + return std::nullopt; + } + + int64_t t_cache_begin = ggml_time_ms(); + if (!copy_cache_tensors_to_cache_buffer(cache_keep_names)) { + if (free_compute_buffer_immediately) { + free_compute_buffer(); + } else if (use_partial_param_offload) { + restore_partial_params(); + } + return std::nullopt; + } + int64_t t_cache_end = ggml_time_ms(); + auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str()); + std::optional> output; + if (!no_return) { + output = sd::make_sd_tensor_from_ggml(result); + } else { + output = sd::Tensor(); + } + + if (free_compute_buffer_immediately) { + free_compute_buffer(); + } else if (use_partial_param_offload) { + restore_partial_params(); + } + if (use_partial_param_offload) { + LOG_DEBUG("%s execute_graph timing: offload=%lld ms alloc=%lld ms copy_in=%lld ms compute=%lld ms cache=%lld ms total=%lld ms", + get_desc().c_str(), + t_offload_end - t_offload_begin, + t_alloc_end - t_alloc_begin, + t_copy_end - t_copy_begin, + t_compute_end - t_compute_begin, + t_cache_end - t_cache_begin, + ggml_time_ms() - t_execute_begin); + } + return output; + } + + template + std::optional> compute_with_graph_cuts(ggml_cgraph* gf, + const GraphCutPlan& plan, + int n_threads, + bool free_compute_buffer_immediately, + bool no_return = false) { + GGML_ASSERT(gf != nullptr); + + free_compute_buffer(); + free_cache_ctx_and_buffer(); + + std::optional> output = sd::Tensor(); + for (size_t seg_idx = 0; seg_idx < plan.segments.size(); ++seg_idx) { + int64_t t_segment_begin = ggml_time_ms(); + const auto& segment = plan.segments[seg_idx]; + auto future_cut_names = sd::ggml_graph_cut::collect_future_input_names(gf, plan, seg_idx); + 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); + if (!bind_segment_cached_inputs(gf, segment)) { + free_cache_ctx_and_buffer(); + free_compute_buffer(); + free_compute_ctx(); + return std::nullopt; + } + + const bool is_last_segment = seg_idx + 1 == plan.segments.size(); + if (!is_last_segment) { + 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); + auto segment_output = execute_graph(segment_graph, + n_threads, + true, + sd::ggml_graph_cut::runtime_param_tensors(gf, segment, get_desc().c_str()), + true, + !is_last_segment || no_return, + &future_cut_names); + 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: virtual std::string get_desc() = 0; GGMLRunner(ggml_backend_t backend, bool offload_params_to_cpu = false) : runtime_backend(backend) { - alloc_params_ctx(); if (!ggml_backend_is_cpu(runtime_backend) && offload_params_to_cpu) { params_backend = ggml_backend_cpu_init(); } else { params_backend = runtime_backend; } + alloc_params_ctx(); } virtual ~GGMLRunner() { @@ -2063,6 +2574,8 @@ struct GGMLRunner { num_tensors); return false; } + rebuild_params_tensor_set(); + ggml_backend_buffer_set_usage(params_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); size_t params_buffer_size = ggml_backend_buffer_get_size(params_buffer); LOG_DEBUG("%s params backend buffer size = % 6.2f MB(%s) (%i tensors)", get_desc().c_str(), @@ -2096,7 +2609,8 @@ struct GGMLRunner { ggml_gallocr_free(compute_allocr); compute_allocr = nullptr; } - offload_params_to_params_backend(); + restore_partial_params(); + restore_all_params(); } // do copy after alloc graph @@ -2160,41 +2674,36 @@ struct GGMLRunner { int n_threads, bool free_compute_buffer_immediately, bool no_return = false) { - if (!offload_params_to_runtime_backend()) { - LOG_ERROR("%s offload params to runtime backend failed", get_desc().c_str()); - return std::nullopt; - } - if (!alloc_compute_buffer(get_graph)) { - LOG_ERROR("%s alloc compute buffer failed", get_desc().c_str()); - return std::nullopt; - } - reset_compute_ctx(); - ggml_cgraph* gf = get_compute_graph(get_graph); - if (!ggml_gallocr_alloc_graph(compute_allocr, gf)) { - LOG_ERROR("%s alloc compute graph failed", get_desc().c_str()); + ggml_cgraph* gf = nullptr; + if (!prepare_compute_graph(get_graph, &gf)) { return std::nullopt; } - copy_data_to_backend_tensor(); - if (ggml_backend_is_cpu(runtime_backend)) { - ggml_backend_cpu_set_n_threads(runtime_backend, n_threads); - } + GGML_ASSERT(gf != nullptr); - ggml_status status = ggml_backend_graph_compute(runtime_backend, gf); - if (status != GGML_STATUS_SUCCESS) { - LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status)); - return std::nullopt; - } - copy_cache_tensors_to_cache_buffer(); - auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str()); - std::optional> output; - if (!no_return) { - output = sd::make_sd_tensor_from_ggml(result); + 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_with_graph_cuts(gf, + plan, + n_threads, + free_compute_buffer_immediately, + no_return); + } } - - if (free_compute_buffer_immediately) { - free_compute_buffer(); + if (!alloc_compute_buffer(gf)) { + LOG_ERROR("%s alloc compute buffer failed", get_desc().c_str()); + return std::nullopt; } - return output; + return execute_graph(gf, + n_threads, + free_compute_buffer_immediately, + {}, + false, + no_return); } void set_flash_attention_enabled(bool enabled) { @@ -2214,6 +2723,10 @@ struct GGMLRunner { weight_adapter = adapter; } + void set_max_graph_vram_bytes(size_t max_vram_bytes) { + max_graph_vram_bytes = max_vram_bytes; + } + ggml_backend_t get_runtime_backend() { return runtime_backend; } diff --git a/src/ggml_graph_cut.cpp b/src/ggml_graph_cut.cpp new file mode 100644 index 000000000..f206f2d2d --- /dev/null +++ b/src/ggml_graph_cut.cpp @@ -0,0 +1,676 @@ +#include "ggml_graph_cut.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "util.h" + +#include "../ggml/src/ggml-impl.h" + +namespace sd::ggml_graph_cut { + + static std::string graph_cut_tensor_display_name(const ggml_tensor* tensor) { + if (tensor == nullptr) { + return ""; + } + if (tensor->name[0] != '\0') { + return tensor->name; + } + return sd_format("", (const void*)tensor); + } + + static int graph_leaf_index(ggml_cgraph* gf, const ggml_tensor* tensor) { + GGML_ASSERT(gf != nullptr); + GGML_ASSERT(tensor != nullptr); + for (int i = 0; i < gf->n_leafs; ++i) { + if (gf->leafs[i] == tensor) { + return i; + } + } + return -1; + } + + static bool is_params_tensor(const std::unordered_set& params_tensor_set, + const ggml_tensor* tensor) { + if (tensor == nullptr) { + return false; + } + return params_tensor_set.find(tensor) != params_tensor_set.end(); + } + + 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 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) { + for (int output_node_index : plan.segments[seg_idx].output_node_indices) { + if (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, + const std::unordered_map& producer_index, + std::unordered_set& available_cut_output_node_indices, + ggml_backend_t backend, + const std::unordered_set& params_tensor_set, + const char* log_desc) { + std::set internal_nodes; + std::unordered_set input_seen; + std::vector input_refs; + + std::stack work_stack; + for (int output_node_index : segment.output_node_indices) { + ggml_tensor* output = ggml_graph_node(gf, output_node_index); + if (output != nullptr) { + work_stack.push(output); + } + } + + while (!work_stack.empty()) { + ggml_tensor* tensor = work_stack.top(); + work_stack.pop(); + + if (tensor == nullptr) { + continue; + } + + auto producer_it = producer_index.find(tensor); + if (producer_it == producer_index.end()) { + if (input_seen.insert(tensor).second) { + Segment::InputRef input_ref; + input_ref.type = is_params_tensor(params_tensor_set, tensor) ? Segment::INPUT_PARAM : Segment::INPUT_EXTERNAL; + input_ref.display_name = graph_cut_tensor_display_name(tensor); + input_ref.leaf_index = graph_leaf_index(gf, tensor); + input_refs.push_back(std::move(input_ref)); + } + continue; + } + + int node_idx = producer_it->second; + if (available_cut_output_node_indices.find(node_idx) != available_cut_output_node_indices.end()) { + if (input_seen.insert(tensor).second) { + Segment::InputRef input_ref; + input_ref.type = Segment::INPUT_PREVIOUS_CUT; + input_ref.display_name = graph_cut_tensor_display_name(tensor); + input_ref.node_index = node_idx; + input_refs.push_back(std::move(input_ref)); + } + continue; + } + + if (!internal_nodes.insert(node_idx).second) { + continue; + } + + ggml_tensor* node = ggml_graph_node(gf, node_idx); + for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { + if (node->src[src_idx] != nullptr) { + work_stack.push(node->src[src_idx]); + } + } + } + + if (!internal_nodes.empty()) { + segment.internal_node_indices.assign(internal_nodes.begin(), internal_nodes.end()); + } + + std::sort(input_refs.begin(), + input_refs.end(), + [](const Segment::InputRef& a, const Segment::InputRef& b) { + if (a.type != b.type) { + return a.type < b.type; + } + 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.compute_buffer_size = measure_segment_compute_buffer(backend, gf, segment, log_desc); + + for (int output_node_index : segment.output_node_indices) { + available_cut_output_node_indices.insert(output_node_index); + } + plan.segments.push_back(std::move(segment)); + } + + bool is_graph_cut_tensor(const ggml_tensor* tensor) { + if (tensor == nullptr || tensor->name[0] == '\0') { + return false; + } + return std::strncmp(tensor->name, GGML_RUNNER_CUT_PREFIX, std::strlen(GGML_RUNNER_CUT_PREFIX)) == 0; + } + + std::string make_graph_cut_name(const std::string& group, const std::string& output) { + return std::string(GGML_RUNNER_CUT_PREFIX) + group + "|" + output; + } + + void mark_graph_cut(ggml_tensor* tensor, const std::string& group, const std::string& output) { + if (tensor == nullptr) { + return; + } + auto name = make_graph_cut_name(group, output); + ggml_set_name(tensor, name.c_str()); + } + + int leaf_count(ggml_cgraph* gf) { + GGML_ASSERT(gf != nullptr); + return gf->n_leafs; + } + + ggml_tensor* leaf_tensor(ggml_cgraph* gf, int leaf_index) { + GGML_ASSERT(gf != nullptr); + if (leaf_index < 0 || leaf_index >= gf->n_leafs) { + return nullptr; + } + return gf->leafs[leaf_index]; + } + + ggml_backend_buffer_t tensor_buffer(const ggml_tensor* tensor) { + if (tensor == nullptr) { + return nullptr; + } + return tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + } + + ggml_tensor* cache_source_tensor(ggml_tensor* tensor) { + if (tensor == nullptr) { + return nullptr; + } + return tensor->view_src ? tensor->view_src : tensor; + } + + size_t cache_tensor_bytes(const ggml_tensor* tensor) { + if (tensor == nullptr) { + return 0; + } + const ggml_tensor* cache_src = tensor->view_src ? tensor->view_src : tensor; + 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; + } + ggml_tensor* leaf = gf->leafs[input_shape_ref.leaf_index]; + if (leaf == nullptr || input_shape_ref.type != leaf->type) { + return false; + } + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (input_shape_ref.ne[static_cast(d)] != leaf->ne[d]) { + return false; + } + } + } + return true; + } + + ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index) { + GGML_ASSERT(gf != nullptr); + if (output_index >= segment.output_node_indices.size()) { + return nullptr; + } + int node_index = segment.output_node_indices[output_index]; + if (node_index < 0 || node_index >= ggml_graph_n_nodes(gf)) { + return nullptr; + } + return ggml_graph_node(gf, node_index); + } + + ggml_tensor* input_tensor(ggml_cgraph* gf, const Segment::InputRef& input_ref) { + GGML_ASSERT(gf != nullptr); + if (input_ref.type == Segment::INPUT_PREVIOUS_CUT) { + if (input_ref.node_index < 0 || input_ref.node_index >= ggml_graph_n_nodes(gf)) { + return nullptr; + } + return ggml_graph_node(gf, input_ref.node_index); + } + if (input_ref.leaf_index < 0 || input_ref.leaf_index >= gf->n_leafs) { + return nullptr; + } + return leaf_tensor(gf, input_ref.leaf_index); + } + + std::vector param_tensors(ggml_cgraph* gf, const Segment& segment) { + GGML_ASSERT(gf != nullptr); + std::vector tensors; + std::unordered_set seen_tensors; + tensors.reserve(segment.input_refs.size()); + seen_tensors.reserve(segment.input_refs.size()); + for (const auto& input_ref : segment.input_refs) { + if (input_ref.type != Segment::INPUT_PARAM) { + continue; + } + ggml_tensor* tensor = input_tensor(gf, input_ref); + if (tensor == nullptr) { + continue; + } + if (seen_tensors.insert(tensor).second) { + tensors.push_back(tensor); + } + } + return tensors; + } + + std::vector runtime_param_tensors(ggml_cgraph* gf, const Segment& segment, const char* log_desc) { + std::vector tensors = param_tensors(gf, segment); + std::vector filtered_tensors; + filtered_tensors.reserve(tensors.size()); + for (ggml_tensor* tensor : tensors) { + if (tensor_buffer(tensor) == nullptr) { + LOG_WARN("%s graph cut skipping param input without buffer: segment=%s tensor=%s", + log_desc == nullptr ? "unknown" : log_desc, + segment.group_name.c_str(), + tensor->name); + continue; + } + filtered_tensors.push_back(tensor); + } + return filtered_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) { + GGML_ASSERT(gf != nullptr); + GGML_ASSERT(graph_ctx_out != nullptr); + + const size_t graph_size = segment.internal_node_indices.size() + segment.input_refs.size() + 8; + ggml_init_params params = { + /*.mem_size =*/ggml_graph_overhead_custom(graph_size, false) + 1024, + /*.mem_buffer =*/nullptr, + /*.no_alloc =*/true, + }; + ggml_context* graph_ctx = ggml_init(params); + GGML_ASSERT(graph_ctx != nullptr); + ggml_cgraph* segment_graph = ggml_new_graph_custom(graph_ctx, graph_size, false); + GGML_ASSERT(segment_graph != nullptr); + + for (const auto& input : segment.input_refs) { + ggml_tensor* current_input = input_tensor(gf, input); + if (current_input == nullptr) { + continue; + } + GGML_ASSERT(segment_graph->n_leafs < segment_graph->size); + segment_graph->leafs[segment_graph->n_leafs++] = current_input; + } + + for (int output_node_index : segment.output_node_indices) { + ggml_tensor* output = ggml_graph_node(gf, output_node_index); + if (output == nullptr) { + continue; + } + ggml_set_output(output); + } + for (int node_idx : segment.internal_node_indices) { + ggml_graph_add_node(segment_graph, ggml_graph_node(gf, node_idx)); + } + *graph_ctx_out = graph_ctx; + return segment_graph; + } + + size_t measure_segment_compute_buffer(ggml_backend_t backend, + ggml_cgraph* gf, + const Segment& segment, + const char* log_desc) { + GGML_ASSERT(backend != nullptr); + GGML_ASSERT(gf != nullptr); + if (segment.internal_node_indices.empty()) { + return 0; + } + + ggml_context* graph_ctx = nullptr; + ggml_cgraph* segment_graph = build_segment_graph(gf, segment, &graph_ctx); + 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, + segment_graph, + nullptr, + nullptr, + sizes); + size_t buffer_size = sizes[0]; + + ggml_gallocr_free(allocr); + ggml_free(graph_ctx); + return buffer_size; + } + + Plan build_plan(ggml_backend_t backend, + ggml_cgraph* gf, + const std::unordered_set& params_tensor_set, + const char* log_desc) { + GGML_ASSERT(backend != nullptr); + GGML_ASSERT(gf != nullptr); + Plan plan; + plan.available = true; + const int n_nodes = ggml_graph_n_nodes(gf); + if (n_nodes <= 0) { + return plan; + } + plan.n_nodes = n_nodes; + plan.n_leafs = gf->n_leafs; + 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); + } + + 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; + } + + std::vector grouped_segments; + std::unordered_map group_to_segment; + for (int i = 0; i < n_nodes; ++i) { + ggml_tensor* node = ggml_graph_node(gf, i); + if (!is_graph_cut_tensor(node)) { + continue; + } + + plan.has_cuts = true; + std::string full_name(node->name); + std::string payload = full_name.substr(std::strlen(GGML_RUNNER_CUT_PREFIX)); + size_t sep = payload.find('|'); + std::string group = sep == std::string::npos ? payload : payload.substr(0, sep); + + auto it = group_to_segment.find(group); + if (it == group_to_segment.end()) { + Segment segment; + segment.group_name = group; + segment.output_node_indices.push_back(i); + group_to_segment[group] = grouped_segments.size(); + grouped_segments.push_back(std::move(segment)); + } else { + auto& segment = grouped_segments[it->second]; + segment.output_node_indices.push_back(i); + } + } + + if (!plan.has_cuts) { + return plan; + } + + std::unordered_set available_cut_output_node_indices; + available_cut_output_node_indices.reserve(static_cast(n_nodes)); + for (auto& segment : grouped_segments) { + build_segment(gf, + plan, + segment, + producer_index, + available_cut_output_node_indices, + backend, + params_tensor_set, + log_desc); + } + + ggml_tensor* final_output = ggml_graph_node(gf, -1); + if (final_output != nullptr && available_cut_output_node_indices.find(n_nodes - 1) == available_cut_output_node_indices.end()) { + Segment final_segment; + final_segment.group_name = "ggml_runner.final"; + final_segment.output_node_indices.push_back(n_nodes - 1); + build_segment(gf, + plan, + final_segment, + producer_index, + available_cut_output_node_indices, + backend, + params_tensor_set, + 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(); + if (graph_cut_segment_vram_bytes(candidate_segment) > max_graph_vram_bytes) { + break; + } + + best_end_segment_index = next_end_segment_index; + } + + 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", + log_desc, + max_graph_vram_bytes / 1024.0 / 1024.0, + base_plan.segments.size(), + merged_plan.segments.size()); + } + + if (log_desc != nullptr) { + LOG_INFO("%s graph cut max_vram budget merge took %lld ms", + log_desc, + ggml_time_ms() - t_budget_begin); + } + + return merged_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); + } + } + + 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; + } + +} // namespace sd::ggml_graph_cut diff --git a/src/ggml_graph_cut.h b/src/ggml_graph_cut.h new file mode 100644 index 000000000..e42859c58 --- /dev/null +++ b/src/ggml_graph_cut.h @@ -0,0 +1,104 @@ +#ifndef __SD_GGML_GRAPH_CUT_H__ +#define __SD_GGML_GRAPH_CUT_H__ + +#include +#include +#include +#include + +#include "ggml-backend.h" +#include "ggml.h" + +namespace sd::ggml_graph_cut { + + struct Segment { + enum InputType { + INPUT_EXTERNAL = 0, + INPUT_PREVIOUS_CUT, + INPUT_PARAM, + }; + + struct InputRef { + InputType type = INPUT_EXTERNAL; + std::string display_name; + int leaf_index = -1; + 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; + std::string group_name; + std::vector internal_node_indices; + std::vector output_node_indices; + std::vector input_refs; + }; + + 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; + 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:"; + + bool is_graph_cut_tensor(const ggml_tensor* tensor); + std::string make_graph_cut_name(const std::string& group, const std::string& output); + void mark_graph_cut(ggml_tensor* tensor, const std::string& group, const std::string& output); + int leaf_count(ggml_cgraph* gf); + ggml_tensor* leaf_tensor(ggml_cgraph* gf, int leaf_index); + 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); + 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::vector runtime_param_tensors(ggml_cgraph* gf, const Segment& segment, const char* log_desc); + 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); + size_t measure_segment_compute_buffer(ggml_backend_t backend, + ggml_cgraph* gf, + const Segment& segment, + const char* log_desc); + Plan build_plan(ggml_backend_t backend, + 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); +} // namespace sd::ggml_graph_cut + +#endif diff --git a/src/llm.hpp b/src/llm.hpp index 4afaa3ba6..a67b4ebf3 100644 --- a/src/llm.hpp +++ b/src/llm.hpp @@ -346,6 +346,7 @@ namespace LLM { auto merger = std::dynamic_pointer_cast(blocks["merger"]); auto x = patch_embed->forward(ctx, pixel_values); + sd::ggml_graph_cut::mark_graph_cut(x, "llm.vision.prelude", "x"); x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0] * spatial_merge_size * spatial_merge_size, x->ne[1] / spatial_merge_size / spatial_merge_size, x->ne[2], x->ne[3]); x = ggml_get_rows(ctx->ggml_ctx, x, window_index); @@ -359,9 +360,11 @@ namespace LLM { mask = nullptr; } x = block->forward(ctx, x, pe, mask); + sd::ggml_graph_cut::mark_graph_cut(x, "llm.vision.blocks." + std::to_string(i), "x"); } x = merger->forward(ctx, x); + sd::ggml_graph_cut::mark_graph_cut(x, "llm.vision.final", "x"); x = ggml_get_rows(ctx->ggml_ctx, x, window_inverse_index); @@ -506,6 +509,7 @@ namespace LLM { auto norm = std::dynamic_pointer_cast(blocks["norm"]); auto x = embed_tokens->forward(ctx, input_ids); + sd::ggml_graph_cut::mark_graph_cut(x, "llm.text.prelude", "x"); std::vector intermediate_outputs; @@ -552,6 +556,10 @@ namespace LLM { auto block = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); x = block->forward(ctx, x, input_pos, attention_mask); + if (out_layers.size() > 1) { + x = ggml_cont(ctx->ggml_ctx, x); + } + sd::ggml_graph_cut::mark_graph_cut(x, "llm.text.layers." + std::to_string(i), "x"); if (out_layers.find(i + 1) != out_layers.end()) { intermediate_outputs.push_back(x); } diff --git a/src/mmdit.hpp b/src/mmdit.hpp index e75736c5d..e57041dc9 100644 --- a/src/mmdit.hpp +++ b/src/mmdit.hpp @@ -767,6 +767,8 @@ struct MMDiT : public GGMLBlock { auto context_x = block->forward(ctx, context, x, c_mod); context = context_x.first; x = context_x.second; + sd::ggml_graph_cut::mark_graph_cut(context, "mmdit.joint_blocks." + std::to_string(i), "context"); + sd::ggml_graph_cut::mark_graph_cut(x, "mmdit.joint_blocks." + std::to_string(i), "x"); } x = final_layer->forward(ctx, x, c_mod); // (N, T, patch_size ** 2 * out_channels) @@ -809,6 +811,11 @@ struct MMDiT : public GGMLBlock { context = context_embedder->forward(ctx, context); // [N, L, D] aka [N, L, 1536] } + sd::ggml_graph_cut::mark_graph_cut(x, "mmdit.prelude", "x"); + sd::ggml_graph_cut::mark_graph_cut(c, "mmdit.prelude", "c"); + if (context != nullptr) { + sd::ggml_graph_cut::mark_graph_cut(context, "mmdit.prelude", "context"); + } x = forward_core_with_concat(ctx, x, c, context, skip_layers); // (N, H*W, patch_size ** 2 * out_channels) diff --git a/src/qwen_image.hpp b/src/qwen_image.hpp index 1cbeb71d5..35d32109e 100644 --- a/src/qwen_image.hpp +++ b/src/qwen_image.hpp @@ -412,6 +412,9 @@ namespace Qwen { auto img = img_in->forward(ctx, x); auto txt = txt_norm->forward(ctx, context); txt = txt_in->forward(ctx, txt); + sd::ggml_graph_cut::mark_graph_cut(img, "qwen_image.prelude", "img"); + sd::ggml_graph_cut::mark_graph_cut(txt, "qwen_image.prelude", "txt"); + // sd::ggml_graph_cut::mark_graph_cut(t_emb, "qwen_image.prelude", "t_emb"); for (int i = 0; i < params.num_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["transformer_blocks." + std::to_string(i)]); @@ -419,6 +422,8 @@ namespace Qwen { auto result = block->forward(ctx, img, txt, t_emb, pe, modulate_index); img = result.first; txt = result.second; + sd::ggml_graph_cut::mark_graph_cut(img, "qwen_image.transformer_blocks." + std::to_string(i), "img"); + sd::ggml_graph_cut::mark_graph_cut(txt, "qwen_image.transformer_blocks." + std::to_string(i), "txt"); } if (params.zero_cond_t) { diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 88102ff61..74c3a07f8 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -144,6 +144,7 @@ class StableDiffusionGGML { std::string taesd_path; sd_tiling_params_t vae_tiling_params = {false, 0, 0, 0.5f, 0, 0}; bool offload_params_to_cpu = false; + float max_vram = 0.f; bool use_pmid = false; bool is_using_v_parameterization = false; @@ -190,6 +191,7 @@ class StableDiffusionGGML { vae_decode_only = sd_ctx_params->vae_decode_only; free_params_immediately = sd_ctx_params->free_params_immediately; offload_params_to_cpu = sd_ctx_params->offload_params_to_cpu; + max_vram = sd_ctx_params->max_vram; bool use_tae = false; @@ -375,6 +377,10 @@ class StableDiffusionGGML { bool clip_on_cpu = sd_ctx_params->keep_clip_on_cpu; + const size_t max_graph_vram_bytes = max_vram <= 0.f + ? 0 + : static_cast(static_cast(max_vram) * 1024.0 * 1024.0 * 1024.0); + { clip_backend = backend; if (clip_on_cpu && !ggml_backend_is_cpu(backend)) { @@ -464,6 +470,7 @@ class StableDiffusionGGML { clip_vision = std::make_shared(backend, offload_params_to_cpu, tensor_storage_map); + clip_vision->set_max_graph_vram_bytes(max_graph_vram_bytes); clip_vision->alloc_params_buffer(); clip_vision->get_param_tensors(tensors); } @@ -540,9 +547,11 @@ class StableDiffusionGGML { } } + cond_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes); cond_stage_model->alloc_params_buffer(); cond_stage_model->get_param_tensors(tensors); + diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes); diffusion_model->alloc_params_buffer(); diffusion_model->get_param_tensors(tensors); @@ -551,6 +560,7 @@ class StableDiffusionGGML { } if (high_noise_diffusion_model) { + high_noise_diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes); high_noise_diffusion_model->alloc_params_buffer(); high_noise_diffusion_model->get_param_tensors(tensors); } @@ -623,16 +633,19 @@ class StableDiffusionGGML { } else if (use_tae && !tae_preview_only) { LOG_INFO("using TAE for encoding / decoding"); first_stage_model = create_tae(); + first_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes); first_stage_model->alloc_params_buffer(); first_stage_model->get_param_tensors(tensors, "tae"); } else { LOG_INFO("using VAE for encoding / decoding"); first_stage_model = create_vae(); + first_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes); first_stage_model->alloc_params_buffer(); first_stage_model->get_param_tensors(tensors, "first_stage_model"); if (use_tae && tae_preview_only) { LOG_INFO("using TAE for preview"); preview_vae = create_tae(); + preview_vae->set_max_graph_vram_bytes(max_graph_vram_bytes); preview_vae->alloc_params_buffer(); preview_vae->get_param_tensors(tensors, "tae"); } @@ -2142,6 +2155,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { sd_ctx_params->prediction = PREDICTION_COUNT; sd_ctx_params->lora_apply_mode = LORA_APPLY_AUTO; sd_ctx_params->offload_params_to_cpu = false; + sd_ctx_params->max_vram = 0.f; sd_ctx_params->enable_mmap = false; sd_ctx_params->keep_clip_on_cpu = false; sd_ctx_params->keep_control_net_on_cpu = false; @@ -2183,6 +2197,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { "sampler_rng_type: %s\n" "prediction: %s\n" "offload_params_to_cpu: %s\n" + "max_vram: %.3f\n" "keep_clip_on_cpu: %s\n" "keep_control_net_on_cpu: %s\n" "keep_vae_on_cpu: %s\n" @@ -2215,6 +2230,7 @@ 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), BOOL_STR(sd_ctx_params->offload_params_to_cpu), + sd_ctx_params->max_vram, BOOL_STR(sd_ctx_params->keep_clip_on_cpu), BOOL_STR(sd_ctx_params->keep_control_net_on_cpu), BOOL_STR(sd_ctx_params->keep_vae_on_cpu), @@ -3432,9 +3448,13 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s std::unique_ptr hires_upscaler; if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) { LOG_INFO("hires fix: loading model upscaler from '%s'", request.hires.model_path); - hires_upscaler = std::make_unique(sd_ctx->sd->n_threads, + hires_upscaler = std::make_unique(sd_ctx->sd->n_threads, false, request.hires.upscale_tile_size); + const size_t max_graph_vram_bytes = sd_ctx->sd->max_vram <= 0.f + ? 0 + : static_cast(static_cast(sd_ctx->sd->max_vram) * 1024.0 * 1024.0 * 1024.0); + hires_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes); if (!hires_upscaler->load_from_file(request.hires.model_path, sd_ctx->sd->offload_params_to_cpu, sd_ctx->sd->n_threads)) { diff --git a/src/t5.hpp b/src/t5.hpp index bbd13e498..71545e522 100644 --- a/src/t5.hpp +++ b/src/t5.hpp @@ -251,7 +251,8 @@ struct T5Stack : public GGMLBlock { ggml_tensor* x, ggml_tensor* past_bias = nullptr, ggml_tensor* attention_mask = nullptr, - ggml_tensor* relative_position_bucket = nullptr) { + ggml_tensor* relative_position_bucket = nullptr, + const std::string& graph_cut_prefix = "") { // x: [N, n_token, model_dim] for (int i = 0; i < num_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["block." + std::to_string(i)]); @@ -259,6 +260,9 @@ struct T5Stack : public GGMLBlock { auto ret = block->forward(ctx, x, past_bias, attention_mask, relative_position_bucket); x = ret.first; past_bias = ret.second; + if (!graph_cut_prefix.empty()) { + sd::ggml_graph_cut::mark_graph_cut(x, graph_cut_prefix + ".block." + std::to_string(i), "x"); + } } auto final_layer_norm = std::dynamic_pointer_cast(blocks["final_layer_norm"]); @@ -305,7 +309,8 @@ struct T5 : public GGMLBlock { auto encoder = std::dynamic_pointer_cast(blocks["encoder"]); auto x = shared->forward(ctx, input_ids); - x = encoder->forward(ctx, x, past_bias, attention_mask, relative_position_bucket); + sd::ggml_graph_cut::mark_graph_cut(x, "t5.prelude", "x"); + x = encoder->forward(ctx, x, past_bias, attention_mask, relative_position_bucket, "t5"); return x; } }; diff --git a/src/unet.hpp b/src/unet.hpp index 2a24f14ed..d7ea8c3fa 100644 --- a/src/unet.hpp +++ b/src/unet.hpp @@ -482,12 +482,14 @@ class UnetModelBlock : public GGMLBlock { emb = ggml_add(ctx->ggml_ctx, emb, label_emb); // [N, time_embed_dim] } + // sd::ggml_graph_cut::mark_graph_cut(emb, "unet.prelude", "emb"); // input_blocks std::vector hs; // input block 0 auto h = input_blocks_0_0->forward(ctx, x); + sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks.0", "h"); ggml_set_name(h, "bench-start"); hs.push_back(h); @@ -505,6 +507,7 @@ class UnetModelBlock : public GGMLBlock { std::string name = "input_blocks." + std::to_string(input_block_idx) + ".1"; h = attention_layer_forward(name, ctx, h, context, num_video_frames); // [N, mult*model_channels, h, w] } + sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks." + std::to_string(input_block_idx), "h"); hs.push_back(h); } if (tiny_unet) { @@ -518,6 +521,7 @@ class UnetModelBlock : public GGMLBlock { auto block = std::dynamic_pointer_cast(blocks[name]); h = block->forward(ctx, h); // [N, mult*model_channels, h/(2^(i+1)), w/(2^(i+1))] + // sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks." + std::to_string(input_block_idx), "h"); hs.push_back(h); } } @@ -531,6 +535,7 @@ class UnetModelBlock : public GGMLBlock { h = resblock_forward("middle_block.2", ctx, h, emb, num_video_frames); // [N, 4*model_channels, h/8, w/8] } } + sd::ggml_graph_cut::mark_graph_cut(h, "unet.middle_block", "h"); if (controls.size() > 0) { auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[controls.size() - 1], control_strength, true); h = ggml_add(ctx->ggml_ctx, h, cs); // middle control @@ -581,6 +586,7 @@ class UnetModelBlock : public GGMLBlock { } output_block_idx += 1; + sd::ggml_graph_cut::mark_graph_cut(h, "unet.output_blocks." + std::to_string(output_block_idx - 1), "h"); } } diff --git a/src/upscaler.cpp b/src/upscaler.cpp index 80e68c947..25fc0c5df 100644 --- a/src/upscaler.cpp +++ b/src/upscaler.cpp @@ -12,6 +12,13 @@ UpscalerGGML::UpscalerGGML(int n_threads, tile_size(tile_size) { } +void UpscalerGGML::set_max_graph_vram_bytes(size_t max_vram_bytes) { + max_graph_vram_bytes = max_vram_bytes; + if (esrgan_upscaler) { + esrgan_upscaler->set_max_graph_vram_bytes(max_vram_bytes); + } +} + bool UpscalerGGML::load_from_file(const std::string& esrgan_path, bool offload_params_to_cpu, int n_threads) { @@ -30,6 +37,7 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path, } LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type)); esrgan_upscaler = std::make_shared(backend, offload_params_to_cpu, tile_size, model_loader.get_tensor_storage_map()); + esrgan_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes); if (direct) { esrgan_upscaler->set_conv2d_direct_enabled(true); } diff --git a/src/upscaler.h b/src/upscaler.h index b11f004a6..d667a6f15 100644 --- a/src/upscaler.h +++ b/src/upscaler.h @@ -14,8 +14,9 @@ struct UpscalerGGML { std::shared_ptr esrgan_upscaler; std::string esrgan_path; int n_threads; - bool direct = false; - int tile_size = 128; + bool direct = false; + int tile_size = 128; + size_t max_graph_vram_bytes = 0; UpscalerGGML(int n_threads, bool direct = false, @@ -24,6 +25,7 @@ struct UpscalerGGML { bool load_from_file(const std::string& esrgan_path, bool offload_params_to_cpu, int n_threads); + void set_max_graph_vram_bytes(size_t max_vram_bytes); 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/wan.hpp b/src/wan.hpp index 6860262c5..261453301 100644 --- a/src/wan.hpp +++ b/src/wan.hpp @@ -692,6 +692,7 @@ namespace WAN { } else { x = conv1->forward(ctx, x); } + // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encoder.prelude", "x"); // downsamples std::vector dims = {dim}; @@ -717,12 +718,14 @@ namespace WAN { x = layer->forward(ctx, x, b, feat_cache, feat_idx, chunk_idx); } } + // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encoder.down." + std::to_string(i), "x"); } // middle x = middle_0->forward(ctx, x, b, feat_cache, feat_idx); x = middle_1->forward(ctx, x, b); x = middle_2->forward(ctx, x, b, feat_cache, feat_idx); + // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encoder.mid", "x"); // head x = head_0->forward(ctx, x); @@ -863,11 +866,13 @@ namespace WAN { } else { x = conv1->forward(ctx, x); } + // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decoder.prelude", "x"); // middle x = middle_0->forward(ctx, x, b, feat_cache, feat_idx); x = middle_1->forward(ctx, x, b); x = middle_2->forward(ctx, x, b, feat_cache, feat_idx); + // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decoder.mid", "x"); // upsamples std::vector dims = {dim_mult[dim_mult.size() - 1] * dim}; @@ -893,6 +898,7 @@ namespace WAN { x = layer->forward(ctx, x, b, feat_cache, feat_idx, chunk_idx); } } + // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decoder.up." + std::to_string(i), "x"); } // head @@ -1031,6 +1037,7 @@ namespace WAN { if (wan2_2) { x = patchify(ctx->ggml_ctx, x, 2, b); } + // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encode.prelude", "x"); auto encoder = std::dynamic_pointer_cast(blocks["encoder"]); auto conv1 = std::dynamic_pointer_cast(blocks["conv1"]); @@ -1051,6 +1058,7 @@ namespace WAN { } out = conv1->forward(ctx, out); auto mu = ggml_ext_chunk(ctx->ggml_ctx, out, 2, 3)[0]; + // sd::ggml_graph_cut::mark_graph_cut(mu, "wan_vae.encode.final", "mu"); clear_cache(); return mu; } @@ -1068,6 +1076,7 @@ namespace WAN { int64_t iter_ = z->ne[2]; auto x = conv2->forward(ctx, z); + // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decode.prelude", "x"); ggml_tensor* out; for (int i = 0; i < iter_; i++) { _conv_idx = 0; @@ -1083,6 +1092,7 @@ namespace WAN { if (wan2_2) { out = unpatchify(ctx->ggml_ctx, out, 2, b); } + // sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode.final", "out"); clear_cache(); return out; } @@ -1097,13 +1107,15 @@ namespace WAN { auto decoder = std::dynamic_pointer_cast(blocks["decoder"]); auto conv2 = std::dynamic_pointer_cast(blocks["conv2"]); - auto x = conv2->forward(ctx, z); + auto x = conv2->forward(ctx, z); + // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decode_partial.prelude", "x"); auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1); // [b*c, 1, h, w] _conv_idx = 0; auto out = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i); if (wan2_2) { out = unpatchify(ctx->ggml_ctx, out, 2, b); } + // sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode_partial.final", "out"); return out; } }; @@ -1984,6 +1996,13 @@ namespace WAN { c = ggml_reshape_3d(ctx->ggml_ctx, c, c->ne[0] * c->ne[1] * c->ne[2], c->ne[3] / N, N); // [N, dim, t_len*h_len*w_len] c = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, c, 1, 0, 2, 3)); // [N, t_len*h_len*w_len, dim] } + sd::ggml_graph_cut::mark_graph_cut(x, "wan.prelude", "x"); + // sd::ggml_graph_cut::mark_graph_cut(e, "wan.prelude", "e"); + // sd::ggml_graph_cut::mark_graph_cut(e0, "wan.prelude", "e0"); + // sd::ggml_graph_cut::mark_graph_cut(context, "wan.prelude", "context"); + if (c != nullptr) { + sd::ggml_graph_cut::mark_graph_cut(c, "wan.prelude", "c"); + } auto x_orig = x; @@ -2004,6 +2023,10 @@ namespace WAN { c_skip = ggml_ext_scale(ctx->ggml_ctx, c_skip, vace_strength); x = ggml_add(ctx->ggml_ctx, x, c_skip); } + sd::ggml_graph_cut::mark_graph_cut(x, "wan.blocks." + std::to_string(i), "x"); + if (c != nullptr) { + sd::ggml_graph_cut::mark_graph_cut(c, "wan.blocks." + std::to_string(i), "c"); + } } x = head->forward(ctx, x, e); // [N, t_len*h_len*w_len, pt*ph*pw*out_dim] diff --git a/src/z_image.hpp b/src/z_image.hpp index 6bb44b791..00b69c264 100644 --- a/src/z_image.hpp +++ b/src/z_image.hpp @@ -371,6 +371,9 @@ namespace ZImage { auto txt = cap_embedder_1->forward(ctx, cap_embedder_0->forward(ctx, context)); // [N, n_txt_token, hidden_size] auto img = x_embedder->forward(ctx, x); // [N, n_img_token, hidden_size] + sd::ggml_graph_cut::mark_graph_cut(txt, "z_image.prelude", "txt"); + sd::ggml_graph_cut::mark_graph_cut(img, "z_image.prelude", "img"); + sd::ggml_graph_cut::mark_graph_cut(t_emb, "z_image.prelude", "t_emb"); int64_t n_txt_pad_token = Rope::bound_mod(static_cast(n_txt_token), SEQ_MULTI_OF); if (n_txt_pad_token > 0) { @@ -393,20 +396,24 @@ namespace ZImage { auto block = std::dynamic_pointer_cast(blocks["context_refiner." + std::to_string(i)]); txt = block->forward(ctx, txt, txt_pe, nullptr, nullptr); + sd::ggml_graph_cut::mark_graph_cut(txt, "z_image.context_refiner." + std::to_string(i), "txt"); } for (int i = 0; i < z_image_params.num_refiner_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["noise_refiner." + std::to_string(i)]); img = block->forward(ctx, img, img_pe, nullptr, t_emb); + sd::ggml_graph_cut::mark_graph_cut(img, "z_image.noise_refiner." + std::to_string(i), "img"); } auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, hidden_size] + sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image.prelude", "txt_img"); for (int i = 0; i < z_image_params.num_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); txt_img = block->forward(ctx, txt_img, pe, nullptr, t_emb); + sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image.layers." + std::to_string(i), "txt_img"); } txt_img = final_layer->forward(ctx, txt_img, t_emb); // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, ph*pw*C]