From 56ba2f14ea44617e94aac75e4926294bdd279064 Mon Sep 17 00:00:00 2001 From: assouan <750048+assouan@users.noreply.github.com> Date: Sat, 22 Aug 2026 01:03:11 +0200 Subject: [PATCH 1/2] fix: cache MiniMax-H3 refined context during sampling Run condition projection and both token-refiner blocks once per conditioning context, then reuse the refined output across denoising steps. Keep entries distinct by condition and active weight adapter. Split each refiner block into its own streaming segment and clear the sampling-scoped cache when sampling finishes. --- src/model/diffusion/minimax_h3.hpp | 145 ++++++++++++++++++++++++++++- src/model/diffusion/model.hpp | 7 ++ src/stable-diffusion.cpp | 21 +++-- 3 files changed, 165 insertions(+), 8 deletions(-) diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index d0683166d..143db3e8d 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -267,11 +268,18 @@ namespace MiniMaxH3 { } ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto final_norm = std::dynamic_pointer_cast(blocks["final_norm"]); for (int64_t i = 0; i < num_layers; ++i) { auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); x = block->forward(ctx, x); + if (i + 1 == num_layers) { + x = final_norm->forward(ctx, x); + } + sd::ggml_graph_cut::mark_graph_cut(x, + "minimax_h3.token_refiner.blocks." + std::to_string(i), + "hidden_states"); } - return std::dynamic_pointer_cast(blocks["final_norm"])->forward(ctx, x); + return num_layers == 0 ? final_norm->forward(ctx, x) : x; } }; @@ -952,6 +960,34 @@ namespace MiniMaxH3 { } struct MiniMaxH3Runner : public DiffusionModelRunner { + struct RefinedContextCacheEntry { + const void* condition_identity = nullptr; + std::shared_ptr weight_adapter = nullptr; + ggml_context* refined_ctx = nullptr; + ggml_backend_buffer_t refined_buffer = nullptr; + ggml_tensor* refined = nullptr; + + ~RefinedContextCacheEntry() { + if (refined_buffer != nullptr) { + ggml_backend_buffer_free(refined_buffer); + } + if (refined_ctx != nullptr) { + ggml_free(refined_ctx); + } + } + + RefinedContextCacheEntry() = default; + RefinedContextCacheEntry(const RefinedContextCacheEntry&) = delete; + RefinedContextCacheEntry& operator=(const RefinedContextCacheEntry&) = delete; + + bool matches(const void* identity, + const std::shared_ptr& adapter) const { + return condition_identity == identity && weight_adapter == adapter; + } + }; + + static constexpr size_t REFINED_CONTEXT_CACHE_CAPACITY = 4; + Config config; MiniMaxH3Transformer3DModel model; sd::Tensor video_input_cache; @@ -961,6 +997,7 @@ namespace MiniMaxH3 { sd::Tensor curve_index_input_cache; sd::Tensor curve_upper_index_input_cache; sd::Tensor curve_fraction_input_cache; + std::vector> refined_context_cache; MiniMaxH3Runner(ggml_backend_t backend, const String2TensorStorage& tensors, @@ -981,6 +1018,92 @@ namespace MiniMaxH3 { model.get_param_tensors(tensors, prefix); } + std::unique_ptr create_refined_context_cache_entry( + const sd::Tensor& context, + const void* condition_identity) { + auto entry = std::make_unique(); + entry->condition_identity = condition_identity; + entry->weight_adapter = weight_adapter; + + auto refined_shape = context.shape(); + refined_shape[0] = config.hidden_size; + ggml_init_params params; + params.mem_size = ggml_tensor_overhead(); + params.mem_buffer = nullptr; + params.no_alloc = true; + entry->refined_ctx = ggml_init(params); + GGML_ASSERT(entry->refined_ctx != nullptr); + entry->refined = ggml_new_tensor(entry->refined_ctx, + GGML_TYPE_F32, + static_cast(refined_shape.size()), + refined_shape.data()); + ggml_set_name(entry->refined, "minimax_h3.refined_context"); + entry->refined_buffer = ggml_backend_alloc_ctx_tensors(entry->refined_ctx, + runtime_backend); + GGML_ASSERT(entry->refined_buffer != nullptr); + ggml_backend_buffer_set_usage(entry->refined_buffer, + GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + return entry; + } + + ggml_cgraph* build_context_refinement_graph(const sd::Tensor& context, + ggml_tensor* refined_output) { + GGML_ASSERT(!context.empty() && context.shape()[0] == config.text_dim); + GGML_ASSERT(refined_output != nullptr && refined_output->ne[0] == config.hidden_size); + auto context_input = make_input(context); + auto runner_ctx = get_context(); + auto refined = model.refine_context(&runner_ctx, context_input); + // Refinement graph buffers are transient; persist only their final output. + auto output = ggml_cpy(runner_ctx.ggml_ctx, refined, refined_output); + auto graph = new_graph_custom(H3_GRAPH_SIZE); + ggml_build_forward_expand(graph, output); + return graph; + } + + ggml_tensor* get_refined_context(const sd::Tensor& context, + const void* condition_identity, + int n_threads) { + GGML_ASSERT(!context.empty()); + GGML_ASSERT(condition_identity != nullptr); + GGML_ASSERT(context.shape()[0] == config.text_dim || + context.shape()[0] == config.hidden_size); + + for (const auto& entry : refined_context_cache) { + if (entry->matches(condition_identity, weight_adapter)) { + return entry->refined; + } + } + + auto entry = create_refined_context_cache_entry(context, condition_identity); + if (context.shape()[0] == config.hidden_size) { + ggml_backend_tensor_set(entry->refined, + context.data(), + 0, + ggml_nbytes(entry->refined)); + ggml_backend_synchronize(runtime_backend); + } else { + auto get_graph = [&]() { + return build_context_refinement_graph(context, entry->refined); + }; + auto result = GGMLRunner::compute(get_graph, + n_threads, + false, + true, + true, + true); + if (!result.has_value()) { + return nullptr; + } + } + + auto refined = entry->refined; + if (refined_context_cache.size() == REFINED_CONTEXT_CACHE_CAPACITY) { + refined_context_cache.erase(refined_context_cache.begin()); + } + refined_context_cache.push_back(std::move(entry)); + return refined; + } + std::pair, sd::Tensor> split_av_latents(const sd::Tensor& packed, int audio_length) const { GGML_ASSERT(packed.dim() == 4 || packed.dim() == 5); @@ -1026,6 +1149,7 @@ namespace MiniMaxH3 { ggml_cgraph* build_graph(const sd::Tensor& packed, const sd::Tensor& timestep, const sd::Tensor& context_tensor, + ggml_tensor* refined_context, const std::vector>& condition_videos, const std::vector>& condition_audios, const sd::Tensor& text_tags, @@ -1039,10 +1163,12 @@ namespace MiniMaxH3 { audio_input_cache = std::move(split.second); GGML_ASSERT(!audio_input_cache.empty()); GGML_ASSERT(!context_tensor.empty()); + GGML_ASSERT(refined_context != nullptr && + refined_context->ne[0] == config.hidden_size); auto video = make_input(video_input_cache); auto audio = make_input(audio_input_cache); - auto context = make_input(context_tensor); + auto context = refined_context; std::vector condition_inputs; condition_inputs.reserve(condition_videos.size()); for (const auto& condition : condition_videos) { @@ -1151,10 +1277,20 @@ namespace MiniMaxH3 { ? empty_reference_blocks : *extra->reference_blocks; const sd::Tensor empty_int; + const void* condition_identity = params.condition_identity != nullptr + ? params.condition_identity + : params.context; + auto context = get_refined_context(*params.context, + condition_identity, + n_threads); + if (context == nullptr) { + return {}; + } auto get_graph = [&]() { return build_graph(*params.x, *params.timesteps, *params.context, + context, conditions, audio_conditions, extra->text_token_tags == nullptr ? empty_int : *extra->text_token_tags, @@ -1171,6 +1307,11 @@ namespace MiniMaxH3 { false), params.x->dim()); } + + protected: + void on_sampling_done() override { + refined_context_cache.clear(); + } }; } // namespace MiniMaxH3 diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index 070ca53d4..f8fde4df0 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -137,6 +137,7 @@ struct DiffusionParams { const sd::Tensor* x = nullptr; const sd::Tensor* timesteps = nullptr; const sd::Tensor* context = nullptr; + const void* condition_identity = nullptr; const sd::Tensor* c_concat = nullptr; const sd::Tensor* y = nullptr; const std::vector>* ref_latents = nullptr; @@ -160,6 +161,7 @@ static inline const sd::Tensor& tensor_or_empty(const sd::Tensor* tensor) struct DiffusionModelRunner : public GGMLRunner { protected: std::string prefix; + virtual void on_sampling_done() {} public: DiffusionModelRunner(ggml_backend_t backend, @@ -171,6 +173,11 @@ struct DiffusionModelRunner : public GGMLRunner { virtual sd::Tensor compute(int n_threads, const DiffusionParams& diffusion_params) = 0; + void sampling_done() { + runner_done(); + on_sampling_done(); + } + void get_param_tensors(std::map& tensors) { get_param_tensors(tensors, prefix); } diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 109a2a483..cf517131e 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -2531,6 +2531,16 @@ class StableDiffusionGGML { float frame_rate, const sd_cache_params_t* cache_params, const sd::Tensor& video_positions = {}) { + struct SamplingDoneOnExit { + DiffusionModelRunner* runner = nullptr; + ~SamplingDoneOnExit() { + if (runner != nullptr) { + runner->sampling_done(); + } + } + }; + SamplingDoneOnExit sample_diffusion_runner_done{work_diffusion_model.get()}; + struct RunnerDoneOnExit { GGMLRunner* runner = nullptr; ~RunnerDoneOnExit() { @@ -2539,8 +2549,6 @@ class StableDiffusionGGML { } } }; - RunnerDoneOnExit sample_diffusion_runner_done{work_diffusion_model.get()}; - RunnerDoneOnExit sample_control_runner_done{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr}; std::vector skip_layers(guidance.slg.layers, guidance.slg.layers + guidance.slg.layer_count); @@ -2712,10 +2720,11 @@ class StableDiffusionGGML { const std::vector* local_skip_layers = nullptr, const std::vector>* ref_latents_override = nullptr, bool use_uncond_ip = false) -> sd::Tensor { - diffusion_params.context = condition.c_crossattn.empty() ? nullptr : &condition.c_crossattn; - diffusion_params.c_concat = c_concat_override != nullptr ? c_concat_override : (condition.c_concat.empty() ? nullptr : &condition.c_concat); - diffusion_params.y = condition.c_vector.empty() ? nullptr : &condition.c_vector; - diffusion_params.ref_latents = ref_latents_override != nullptr ? ref_latents_override : (condition.c_ref_images.empty() ? &ref_latents : &condition.c_ref_images); + diffusion_params.context = condition.c_crossattn.empty() ? nullptr : &condition.c_crossattn; + diffusion_params.condition_identity = &condition; + diffusion_params.c_concat = c_concat_override != nullptr ? c_concat_override : (condition.c_concat.empty() ? nullptr : &condition.c_concat); + diffusion_params.y = condition.c_vector.empty() ? nullptr : &condition.c_vector; + diffusion_params.ref_latents = ref_latents_override != nullptr ? ref_latents_override : (condition.c_ref_images.empty() ? &ref_latents : &condition.c_ref_images); if (sd_version_is_unet(version)) { int nvf = -1; From 34a3146c6948f4cd0039af9f1ed427c7308947ef Mon Sep 17 00:00:00 2001 From: assouan <750048+assouan@users.noreply.github.com> Date: Sat, 22 Aug 2026 02:48:45 +0200 Subject: [PATCH 2/2] fix: harden MiniMax-H3 refined context streaming Split condition projection and token refiners into sequential streamed segments while folding the persistent copy into the final refiner segment. Tie cached refined contexts to the source tensor identity, storage, shape, and active weight adapter. --- src/model/diffusion/minimax_h3.hpp | 70 ++++++++++++++++++++---------- src/model/diffusion/model.hpp | 2 +- src/stable-diffusion.cpp | 10 ++--- 3 files changed, 52 insertions(+), 30 deletions(-) diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 143db3e8d..c3b9c0690 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -267,17 +267,22 @@ namespace MiniMaxH3 { blocks["final_norm"] = std::make_shared(config.hidden_size, config.final_norm_eps); } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + bool cut_after_last = true) { auto final_norm = std::dynamic_pointer_cast(blocks["final_norm"]); for (int64_t i = 0; i < num_layers; ++i) { auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); x = block->forward(ctx, x); - if (i + 1 == num_layers) { + const bool is_last = i + 1 == num_layers; + if (is_last) { x = final_norm->forward(ctx, x); } - sd::ggml_graph_cut::mark_graph_cut(x, - "minimax_h3.token_refiner.blocks." + std::to_string(i), - "hidden_states"); + if (!is_last || cut_after_last) { + sd::ggml_graph_cut::mark_graph_cut(x, + "minimax_h3.token_refiner.blocks." + std::to_string(i), + "hidden_states"); + } } return num_layers == 0 ? final_norm->forward(ctx, x) : x; } @@ -534,14 +539,20 @@ namespace MiniMaxH3 { } } - ggml_tensor* refine_context(GGMLRunnerContext* ctx, ggml_tensor* context) { + ggml_tensor* refine_context(GGMLRunnerContext* ctx, + ggml_tensor* context, + bool cut_after_last_refiner = true) { if (context->ne[0] == config.hidden_size) { return context; } GGML_ASSERT(context->ne[0] == config.text_dim); auto condition_proj = std::dynamic_pointer_cast(blocks["condition_proj"]); auto token_refiner = std::dynamic_pointer_cast(blocks["token_refiner"]); - return token_refiner->forward(ctx, condition_proj->forward(ctx, context)); + auto projected = condition_proj->forward(ctx, context); + sd::ggml_graph_cut::mark_graph_cut(projected, + "minimax_h3.condition_proj", + "hidden_states"); + return token_refiner->forward(ctx, projected, cut_after_last_refiner); } ggml_tensor* time_embedding(GGMLRunnerContext* ctx, @@ -961,7 +972,10 @@ namespace MiniMaxH3 { struct MiniMaxH3Runner : public DiffusionModelRunner { struct RefinedContextCacheEntry { - const void* condition_identity = nullptr; + const void* context_cache_identity = nullptr; + const sd::Tensor* source_context = nullptr; + const float* source_data = nullptr; + std::vector source_shape; std::shared_ptr weight_adapter = nullptr; ggml_context* refined_ctx = nullptr; ggml_backend_buffer_t refined_buffer = nullptr; @@ -981,8 +995,13 @@ namespace MiniMaxH3 { RefinedContextCacheEntry& operator=(const RefinedContextCacheEntry&) = delete; bool matches(const void* identity, + const sd::Tensor& context, const std::shared_ptr& adapter) const { - return condition_identity == identity && weight_adapter == adapter; + return context_cache_identity == identity && + weight_adapter == adapter && + source_context == &context && + source_data == context.data() && + source_shape == context.shape(); } }; @@ -1020,10 +1039,13 @@ namespace MiniMaxH3 { std::unique_ptr create_refined_context_cache_entry( const sd::Tensor& context, - const void* condition_identity) { - auto entry = std::make_unique(); - entry->condition_identity = condition_identity; - entry->weight_adapter = weight_adapter; + const void* context_cache_identity) { + auto entry = std::make_unique(); + entry->context_cache_identity = context_cache_identity; + entry->source_context = &context; + entry->source_data = context.data(); + entry->source_shape = context.shape(); + entry->weight_adapter = weight_adapter; auto refined_shape = context.shape(); refined_shape[0] = config.hidden_size; @@ -1052,7 +1074,7 @@ namespace MiniMaxH3 { GGML_ASSERT(refined_output != nullptr && refined_output->ne[0] == config.hidden_size); auto context_input = make_input(context); auto runner_ctx = get_context(); - auto refined = model.refine_context(&runner_ctx, context_input); + auto refined = model.refine_context(&runner_ctx, context_input, false); // Refinement graph buffers are transient; persist only their final output. auto output = ggml_cpy(runner_ctx.ggml_ctx, refined, refined_output); auto graph = new_graph_custom(H3_GRAPH_SIZE); @@ -1061,20 +1083,20 @@ namespace MiniMaxH3 { } ggml_tensor* get_refined_context(const sd::Tensor& context, - const void* condition_identity, + const void* context_cache_identity, int n_threads) { GGML_ASSERT(!context.empty()); - GGML_ASSERT(condition_identity != nullptr); + GGML_ASSERT(context_cache_identity != nullptr); GGML_ASSERT(context.shape()[0] == config.text_dim || context.shape()[0] == config.hidden_size); for (const auto& entry : refined_context_cache) { - if (entry->matches(condition_identity, weight_adapter)) { + if (entry->matches(context_cache_identity, context, weight_adapter)) { return entry->refined; } } - auto entry = create_refined_context_cache_entry(context, condition_identity); + auto entry = create_refined_context_cache_entry(context, context_cache_identity); if (context.shape()[0] == config.hidden_size) { ggml_backend_tensor_set(entry->refined, context.data(), @@ -1277,12 +1299,12 @@ namespace MiniMaxH3 { ? empty_reference_blocks : *extra->reference_blocks; const sd::Tensor empty_int; - const void* condition_identity = params.condition_identity != nullptr - ? params.condition_identity - : params.context; - auto context = get_refined_context(*params.context, - condition_identity, - n_threads); + const void* context_cache_identity = params.context_cache_identity != nullptr + ? params.context_cache_identity + : params.context; + auto context = get_refined_context(*params.context, + context_cache_identity, + n_threads); if (context == nullptr) { return {}; } diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index f8fde4df0..f895e5eae 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -137,7 +137,7 @@ struct DiffusionParams { const sd::Tensor* x = nullptr; const sd::Tensor* timesteps = nullptr; const sd::Tensor* context = nullptr; - const void* condition_identity = nullptr; + const void* context_cache_identity = nullptr; const sd::Tensor* c_concat = nullptr; const sd::Tensor* y = nullptr; const std::vector>* ref_latents = nullptr; diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index cf517131e..c7587b0e3 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -2720,11 +2720,11 @@ class StableDiffusionGGML { const std::vector* local_skip_layers = nullptr, const std::vector>* ref_latents_override = nullptr, bool use_uncond_ip = false) -> sd::Tensor { - diffusion_params.context = condition.c_crossattn.empty() ? nullptr : &condition.c_crossattn; - diffusion_params.condition_identity = &condition; - diffusion_params.c_concat = c_concat_override != nullptr ? c_concat_override : (condition.c_concat.empty() ? nullptr : &condition.c_concat); - diffusion_params.y = condition.c_vector.empty() ? nullptr : &condition.c_vector; - diffusion_params.ref_latents = ref_latents_override != nullptr ? ref_latents_override : (condition.c_ref_images.empty() ? &ref_latents : &condition.c_ref_images); + diffusion_params.context = condition.c_crossattn.empty() ? nullptr : &condition.c_crossattn; + diffusion_params.context_cache_identity = &condition; + diffusion_params.c_concat = c_concat_override != nullptr ? c_concat_override : (condition.c_concat.empty() ? nullptr : &condition.c_concat); + diffusion_params.y = condition.c_vector.empty() ? nullptr : &condition.c_vector; + diffusion_params.ref_latents = ref_latents_override != nullptr ? ref_latents_override : (condition.c_ref_images.empty() ? &ref_latents : &condition.c_ref_images); if (sd_version_is_unet(version)) { int nvf = -1;