From 5d0ad3e31212b20edd21974f02f444e9a219d01b Mon Sep 17 00:00:00 2001 From: Liwei Lan Date: Wed, 15 Jul 2026 14:15:37 +0800 Subject: [PATCH 1/8] refactor: delete the CacheGraphScope seam; tap registry is the only cache path --- src/core/backend/ggml/ggml_extend.hpp | 193 ++------- .../optimization/cache/cache_graph_scope.hpp | 208 +-------- .../cache/compile/cache_graph_lowering.cpp | 144 +++---- .../cache/compile/cache_graph_lowering.hpp | 61 +-- .../optimization/cache/ir/substep_plan.hpp | 2 + .../optimization/cache/model/tap_registry.hpp | 19 + .../cache/policy/cache_policy.hpp | 15 +- .../policy/policies/condition_policy.cpp | 8 +- .../cache/policy/policies/dicache_policy.cpp | 46 +- .../policy/policies/easycache_policy.cpp | 9 +- .../cache/policy/policies/magcache_policy.cpp | 14 +- .../cache/policy/policies/sencache_policy.cpp | 9 +- .../policy/policies/taylorseer_policy.cpp | 8 +- .../cache/policy/policies/ucache_policy.cpp | 8 +- .../cache/runtime/cache_engine.cpp | 26 +- .../cache/state/cache_device_store.hpp | 2 +- src/dit_models/diffusion_model.hpp | 120 +----- src/dit_models/models/flux.hpp | 398 ++++-------------- src/dit_models/models/mmdit.hpp | 93 +--- src/dit_models/models/qwen_image.hpp | 337 +++------------ src/dit_models/models/wan.hpp | 106 +---- .../pipelines/flux_kontext_pipeline.cpp | 23 - src/dit_models/pipelines/flux_pipeline.cpp | 54 +-- .../pipelines/qwen_image_edit_pipeline.cpp | 59 --- .../pipelines/qwen_image_pipeline.cpp | 48 +-- src/dit_models/pipelines/sd3_pipeline.cpp | 25 +- src/dit_models/pipelines/wan_pipeline.cpp | 19 +- 27 files changed, 387 insertions(+), 1667 deletions(-) diff --git a/src/core/backend/ggml/ggml_extend.hpp b/src/core/backend/ggml/ggml_extend.hpp index ae791fb2..43cd3958 100644 --- a/src/core/backend/ggml/ggml_extend.hpp +++ b/src/core/backend/ggml/ggml_extend.hpp @@ -1808,9 +1808,7 @@ struct GGMLRunnerContext { bool circular_x_enabled = false; bool circular_y_enabled = false; std::shared_ptr weight_adapter = nullptr; - // Build-time cache seam; null on the uncached path (graph is then identical). - sd::CacheGraphScope* cache_scope = nullptr; - // Substep tap registry; null on the legacy/uncached path (no taps recorded, + // Substep tap registry; null on the uncached path (no taps recorded, // graph identical). When set, the model's forward() conditionally taps // structural anchors it requests (see cache/model/tap_registry.hpp). edgedit::cache::TapRegistry* tap_registry = nullptr; @@ -1835,8 +1833,8 @@ inline void tap(GGMLRunnerContext* ctx, const edgedit::cache::AnchorRef& a, ggml // Reconstruct the reuse output for a tap-driven inject: x_before + , // where the residual depends on the registry's inject kind (host feature / device -// single residual / device gamma-blend). Mirrors CacheGraphScope::add_injected but -// registry-driven. Called by a model forward at the inject region start. +// single residual / device gamma-blend). Called by a model forward at the inject +// region start. inline ggml_tensor* build_tap_inject(GGMLRunnerContext* ctx, ggml_tensor* x_before) { using edgedit::cache::TapRegistry; ggml_context* c = ctx->ggml_ctx; @@ -1861,7 +1859,7 @@ inline ggml_tensor* build_tap_inject(GGMLRunnerContext* ctx, ggml_tensor* x_befo // Device backing for CacheStateManager slots (implements the cache core's // abstract ICacheDeviceStore). Owns one persistent ggml_context + backend buffer, // allocated on the runtime backend OUTSIDE any per-step compute buffer, so slot -// tensors survive run_cache_pass's compute(free=false) + reset_graph_cut_run_cache(). +// tensors survive the substep pass's compute(free=false) + reset_graph_cut_run_cache(). // Entries are keyed by (ring_key, ring_index); a shape change frees and rebuilds // (mirrors ensure_dicache_gpu_state). release_all() runs per generation from // CacheStateManager::reset(), matching the old reset_dicache_gpu_states() lifetime. @@ -2296,13 +2294,8 @@ struct GGMLRunner { bool circular_x_enabled = false; bool circular_y_enabled = false; - // Non-owning; set by the CacheController before a cache pass, cleared after. - // When null, get_context() leaves ctx.cache_scope null and the model graph - // is byte-identical to the uncached path. - sd::CacheGraphScope* cache_scope_ = nullptr; - // Non-owning substep tap registry; set before a substep-path build, cleared - // after. Null on legacy/uncached builds (graph identical). Threaded onto the + // after. Null on uncached builds (graph identical). Threaded onto the // GGMLRunnerContext so the model's tap() calls land here. edgedit::cache::TapRegistry* tap_registry_ = nullptr; @@ -4668,9 +4661,6 @@ struct GGMLRunner { auto result = ggml_graph_node(gf, -1); ggml_set_name(result, final_result_name.c_str()); } - if (cache_scope_ != nullptr && ggml_graph_n_nodes(gf) > 0) { - expand_cache_scope_nodes(gf); - } if (tap_registry_ != nullptr && ggml_graph_n_nodes(gf) > 0) { expand_tap_registry_nodes(gf); } @@ -4679,12 +4669,10 @@ struct GGMLRunner { } // Promote the substep's recorded taps + woven indicator scalars into the - // named-tensor index so they can be read back post-compute (mirrors - // expand_cache_scope_nodes, but the set is dynamic and plan-driven rather than - // fixed *_node fields). Each tapped anchor is pinned as a graph output under a - // stable name; indicator scalar nodes name themselves (cache_ind:) and - // are expanded here too. The registry only holds anchors the model actually - // tapped this build. + // named-tensor index so they can be read back post-compute. Each tapped anchor + // is pinned as a graph output under a stable name; indicator scalar nodes name + // themselves (cache_ind:) and are expanded here too. The registry only + // holds anchors the model actually tapped this build. void expand_tap_registry_nodes(ggml_cgraph* gf) { auto expand_named = [&](ggml_tensor* node, const std::string& name) { if (node == nullptr) { @@ -4697,8 +4685,7 @@ struct GGMLRunner { }; // Weave indicators + residual capture HERE (after the final-result node was // named at prepare-time), so the model output stays the graph's result and - // these aux nodes are appended without stealing the result slot — mirrors - // expand_cache_scope_nodes' ordering. + // these aux nodes are appended without stealing the result slot. for (const auto& ind : tap_registry_->indicators()) { ggml_tensor* s = edgedit::cache::lower_indicator(compute_ctx, ind, *tap_registry_); if (s != nullptr) { @@ -4713,6 +4700,20 @@ struct GGMLRunner { expand_named(feat, "ed_cache_feature"); } } + // Capture-writeback (DiCache device seed): also weave the probe residual + // (BlockOut[m-1] - ModelIn) as "ed_cache_probe_resid" so the pass's + // post-readback can d2d it into the runner's persistent probe-residual ring. + // The raw probe/before taps are already pinned by the recorded() loop below, + // so the handoff reads those by their ed_tap: names. Same sub-expr form as the + // gamma weave above. + if (tap_registry_->capture_writeback()) { + const int m = tap_registry_->writeback_probe_depth(); + ggml_tensor* before = tap_registry_->get(edgedit::cache::AnchorRef::model_in()); + ggml_tensor* probe = tap_registry_->get(edgedit::cache::AnchorRef::block_out(m - 1)); + if (before != nullptr && probe != nullptr) { + expand_named(ggml_sub(compute_ctx, probe, before), "ed_cache_probe_resid"); + } + } // DiCache probe metrics (delta_y/delta_x/gamma): weave the exact // build_probe_metrics form from the tapped probe (BlockOut[m]) + before // (ModelIn) anchors and the runner-owned persistent operands. Only the @@ -4750,46 +4751,6 @@ struct GGMLRunner { } } - // Names used to read back the cache seam's aux tensors post-compute. - static constexpr const char* kCacheFeatureName = "ed_cache_feature"; - static constexpr const char* kCacheBeforeName = "ed_cache_before"; - static constexpr const char* kCacheProbeName = "ed_cache_probe"; - static constexpr const char* kCacheDeltaYName = "ed_cache_delta_y"; - static constexpr const char* kCacheDeltaXName = "ed_cache_delta_x"; - static constexpr const char* kCacheGammaName = "ed_cache_gamma"; - static constexpr const char* kCacheProbeResidName = "ed_cache_probe_resid"; - - void expand_cache_scope_nodes(ggml_cgraph* gf) { - auto expand_named = [&](ggml_tensor* node, const char* name) { - if (node == nullptr) { - return; - } - ggml_set_name(node, name); - ggml_set_output(node); // keep the allocator from reusing its buffer - cache(name, node); - ggml_build_forward_expand(gf, node); - }; - if (cache_scope_->capture_mode()) { - expand_named(cache_scope_->feature_node, kCacheFeatureName); - // GPU DiCache: also expose before/probe/probe_resid so capture can - // snapshot them into the persistent cross-step buffers (device-to-device). - expand_named(cache_scope_->before_node, kCacheBeforeName); - expand_named(cache_scope_->probe_node, kCacheProbeName); - expand_named(cache_scope_->probe_resid_node, kCacheProbeResidName); - } else if (cache_scope_->probe_mode()) { - // On probe the model returns the probe state as `out`, so it becomes - // the final result node; only the block-stack input needs a name. - expand_named(cache_scope_->before_node, kCacheBeforeName); - expand_named(cache_scope_->probe_node, kCacheProbeName); - // GPU DiCache: the decision scalars (built from probe/before vs the - // persistent prev_* tensors). Only a few bytes are read back per step. - cache_scope_->build_probe_metrics(compute_ctx); - expand_named(cache_scope_->delta_y_node, kCacheDeltaYName); - expand_named(cache_scope_->delta_x_node, kCacheDeltaXName); - expand_named(cache_scope_->gamma_node, kCacheGammaName); - } - } - bool prepare_compute_graph(get_graph_cb_t get_graph, ggml_cgraph** gf_out) { GGML_ASSERT(gf_out != nullptr); @@ -6071,7 +6032,6 @@ struct GGMLRunner { runner_ctx.circular_x_enabled = circular_x_enabled; runner_ctx.circular_y_enabled = circular_y_enabled; runner_ctx.weight_adapter = weight_adapter; - runner_ctx.cache_scope = cache_scope_; runner_ctx.tap_registry = tap_registry_; return runner_ctx; } @@ -6263,102 +6223,7 @@ struct GGMLRunner { return true; } - // Run one cache-aware compute pass with `scope` attached. Builds the graph - // via `get_graph` (whose model forward() consults the scope), executes on - // the plain path (feature caching is gated off the segmented path by the - // pipeline), and reads the seam's aux tensors back to host. `expected_dim` - // restores trailing singleton dims on the main result like compute() does. - // The compute buffer is kept alive (not freed) so the named aux nodes are - // readable; the caller frees it as usual after the step. - struct CachePassResult { - sd::Tensor output; - sd::Tensor feature; - sd::Tensor before; - sd::Tensor probe; - // GPU DiCache scalar readbacks (NaN if the node was absent this pass). - float delta_y = std::numeric_limits::quiet_NaN(); - float delta_x = std::numeric_limits::quiet_NaN(); - float gamma = std::numeric_limits::quiet_NaN(); - }; - CachePassResult run_cache_pass(get_graph_cb_t get_graph, - int n_threads, - sd::CacheGraphScope* scope, - size_t expected_dim, - const std::function& post_readback = nullptr) { - CachePassResult result; - set_cache_scope(scope); - const int64_t t_pass_begin = ggml_time_ms(); - const bool dprof = std::getenv("ED_PROFILE_DICACHE") != nullptr; - auto out = GGMLRunner::compute(get_graph, n_threads, /*free=*/false); - result.output = restore_trailing_singleton_dims(std::move(out), expected_dim); - if (dprof && scope != nullptr) { - const char* m = scope->probe_mode() ? "probe" - : scope->inject_mode() ? "inject" - : scope->capture_mode() ? "capture" : "other"; - LOG_INFO("[dicache-prof] %s pass %lld ms", m, - (long long)(ggml_time_ms() - t_pass_begin)); - } - - if (scope != nullptr) { - auto read_scalar = [&](const char* name) -> float { - ggml_tensor* t = get_cache_tensor_by_name(name); - if (t == nullptr) { - return std::numeric_limits::quiet_NaN(); - } - float v = std::numeric_limits::quiet_NaN(); - ggml_backend_tensor_get(t, &v, 0, sizeof(float)); - return v; - }; - if (scope->capture_mode()) { - // Host path only: read the ~50-190MB feature residual back so the - // policy can build its host residual ring in observe(). GPU DiCache - // fills that ring device-to-device in the post_readback handoff and - // reconstructs via inject_gpu, so the host readback is dead weight — - // skip it under gpu_metric. - if (!scope->gpu_metric) { - ggml_tensor* feat = get_cache_tensor_by_name(kCacheFeatureName); - if (feat != nullptr) { - result.feature = sd::make_sd_tensor_from_ggml(feat); - } - } - } else if (scope->probe_mode()) { - if (scope->gpu_metric) { - // GPU DiCache: read back only the decision scalars (a few bytes), - // NOT the ~50MB before/probe tensors. - result.delta_y = read_scalar(kCacheDeltaYName); - result.delta_x = read_scalar(kCacheDeltaXName); - result.gamma = read_scalar(kCacheGammaName); - } else { - ggml_tensor* before = get_cache_tensor_by_name(kCacheBeforeName); - if (before != nullptr) { - result.before = sd::make_sd_tensor_from_ggml(before); - } - ggml_tensor* probe = get_cache_tensor_by_name(kCacheProbeName); - if (probe != nullptr) { - result.probe = sd::make_sd_tensor_from_ggml(probe); - } - } - } - } - set_cache_scope(nullptr); - // GPU DiCache handoff: while the named cache tensors are still resident - // (before reset frees the run cache), copy them device-to-device into the - // runner's persistent cross-step buffers. Callback uses get_cache_tensor_by_name. - if (post_readback) { - post_readback(); - } - // compute(free=false) leaves the run cache (cache_ctx/buffer/chunks) - // allocated; the plain (non-segmented) seam path never calls - // reset_graph_cut_run_cache(), so without this a per-step chunk would - // survive across steps AND generations (~1GB/img GPU leak). Free it here - // now that the host copies are made; pool retention keeps one same-layout - // chunk for cheap reuse. - reset_graph_cut_run_cache(); - return result; - } - - // ---- Substep-path pass (ED_CACHE_SUBSTEP). Tap-driven analogue of - // run_cache_pass: the middle layer configures a TapRegistry (requested anchors, + // ---- Substep-path pass. The middle layer configures a TapRegistry (requested anchors, // indicators, stop-after-block) before calling; the model forward() taps the // anchors and build_graph weaves the indicator scalars. This reads back the // requested indicator scalars by name and, when a residual capture is asked for, @@ -6594,19 +6459,13 @@ struct GGMLRunner { flash_attn_enabled = enabled; } - // Attach/detach the build-time cache seam consulted by model forward(). - void set_cache_scope(sd::CacheGraphScope* scope) { - cache_scope_ = scope; - } + // Attach/detach the substep tap registry consulted by model forward(). void set_tap_registry(edgedit::cache::TapRegistry* reg) { tap_registry_ = reg; } edgedit::cache::TapRegistry* tap_registry() const { return tap_registry_; } - sd::CacheGraphScope* cache_scope() const { - return cache_scope_; - } // True when the block-stack feature seam can run: only on the plain compute // path (no process-group comm, no VRAM-budgeted segmented execution), since // mid-graph capture is not preserved across segments. diff --git a/src/core/optimization/cache/cache_graph_scope.hpp b/src/core/optimization/cache/cache_graph_scope.hpp index f276893e..00432b3b 100644 --- a/src/core/optimization/cache/cache_graph_scope.hpp +++ b/src/core/optimization/cache/cache_graph_scope.hpp @@ -1,17 +1,16 @@ #pragma once -#include #include -#include "ggml.h" #include "utils/tensor.hpp" namespace sd { -// Result of a cache-aware compute pass, produced by the runner device helpers -// and consumed by the CacheController. `output` is the usual denoiser output +// Result of a cache-aware compute pass, produced by the runner's substep helpers +// and consumed by the cache lowering. `output` is the usual denoiser output // (empty on a probe pass). `feature`/`before`/`probe` carry the block-stack -// tensors a Feature/Probe policy needs, already read back to host. +// tensors a Feature/Probe policy needs, already read back to host; the scalar +// fields carry the on-device DiCache decision metrics (NaN when not produced). struct DiffusionCacheResult { Tensor output; Tensor feature; // Capture: cached-region residual @@ -23,203 +22,4 @@ struct DiffusionCacheResult { float gamma = std::numeric_limits::quiet_NaN(); }; -// Build-time cache seam shared by all DiT model graphs. A model's forward() -// consults a nullable CacheGraphScope* (carried on GGMLRunnerContext) inside its -// block loop; when null the graph is byte-identical to the uncached path. The -// scope is method-AGNOSTIC — it only knows the graph shapes a cache step can -// take, never which cache algorithm requested them. The runner owns the -// ggml_cgraph and does the cache()/ggml_set_output/expand wiring using the -// tensor pointers the model records here. -// -// The cached region is the contiguous block interval [region_start, region_end) -// of the model's block stack. region_end < 0 means "to the end of the stack". -// The default (0, -1) is the whole stack, which the current whole-model methods -// (TaylorSeer / MagCache / DiCache) use. A policy MAY instead narrow it to a -// sub-interval: blocks before region_start and from region_end on always run; -// only the region is captured on full steps and skipped (residual injected) on -// reuse steps. This costs nothing beyond the one compute pass the step already -// runs — it is a pure build-time structural decision, not an extra pass. -// -// Modes: -// Capture - run every block; expose feature = region_after - region_before. -// Inject - run the blocks OUTSIDE the region; at region_start, replace the -// region with region_before + inject_feature and resume at region_end. -// Probe - run probe_depth blocks from region_start; expose that hidden state, -// then the model returns it directly (final layer is not run). -// -// Two API levels: -// * In-loop sugar (begin_region / step_inject_region / end_region / -// stop_after_block) — used by uniform single-loop models (mmdit, qwen, wan) -// and by flux for sub-regions inside its double-block loop. -// * Low-level primitives (snapshot_before / build_feature / add_injected) — -// used by flux's whole-stack path, whose residual spans the double AND single -// loops and is captured after the streams are recombined. -struct CacheGraphScope { - enum class Mode { Capture, Inject, Probe }; - - Mode mode = Mode::Capture; - int probe_depth = 0; - - // Cached region = blocks [region_start, region_end). Defaults to whole stack. - int region_start = 0; - int region_end = -1; // < 0 => to the end of the stack - - // Inject input: reconstructed residual, created by the runner with the same - // shape as the region's input state. - ggml_tensor* inject_feature = nullptr; - - // ---- GPU-side DiCache (Probe granularity) ---- - // When these persistent cross-step tensors are provided, the seam computes the - // DiCache decision metric and residual reconstruction ON the GPU instead of - // reading ~50MB tensors back to host every step. All are runner-owned and - // persist across the sampling loop; nullptr => host-side path (or non-DiCache). - ggml_tensor* prev_probe_node = nullptr; // probe state from the last computed step - ggml_tensor* prev_input_node = nullptr; // block-stack input from last computed step - ggml_tensor* probe_prev1_node = nullptr; // newest probe residual (probe - before) - ggml_tensor* probe_prev2_node = nullptr; // 2nd-newest probe residual - ggml_tensor* resid_prev1_node = nullptr; // newest full block-stack residual - ggml_tensor* resid_prev2_node = nullptr; // 2nd-newest full residual - ggml_tensor* gamma_scalar = nullptr; // Inject: trajectory-alignment gamma [1] - bool gpu_metric = false; // Probe: emit delta_y/delta_x/gamma scalar nodes - - // Recorded by the runner after a Probe pass (scalar readbacks); consumed by the - // policy. Named nodes: kCacheDeltaYName / kCacheDeltaXName / kCacheGammaName. - ggml_tensor* delta_y_node = nullptr; - ggml_tensor* delta_x_node = nullptr; - ggml_tensor* gamma_node = nullptr; - ggml_tensor* probe_resid_node = nullptr; // Capture: (probe_state - block_input) - - // Recorded by the model, consumed by the runner after forward() returns. - ggml_tensor* before_node = nullptr; // region input (Capture anchor / Probe readback) - ggml_tensor* probe_node = nullptr; // hidden state after probe_depth blocks (Probe) - ggml_tensor* feature_node = nullptr; // region_after - region_before (Capture) - - bool capture_mode() const { return mode == Mode::Capture; } - bool inject_mode() const { return mode == Mode::Inject; } - bool probe_mode() const { return mode == Mode::Probe; } - - // A finite region_end marks a sub-interval; region_end < 0 is the whole stack. - bool is_subregion() const { return region_end >= 0; } - int resolved_region_end(int n) const { return region_end < 0 ? n : std::min(region_end, n); } - bool at_region_start(int i) const { return i == region_start; } - // Block index to resume at after an injected region is skipped. - int inject_resume_index(int n) const { return resolved_region_end(n); } - - // ---- Low-level primitives (whole-stack path) ---- - void snapshot_before(ggml_tensor* x) { before_node = x; } - ggml_tensor* add_injected(ggml_context* ctx, ggml_tensor* x_before) const { - // GPU-side Dynamic Cache Trajectory Alignment: when the persistent residual - // nodes and the gamma scalar are present, reconstruct - // x_before + resid_prev2 + gamma*(resid_prev1 - resid_prev2) - // entirely on-device (matches ref _dicache_apply_cached_residual), instead - // of injecting a host-reconstructed residual. Falls back to the simple - // x_before + inject_feature path for host-side / feature-cache methods. - if (gpu_reconstruct_available()) { - ggml_tensor* diff = ggml_sub(ctx, resid_prev1_node, resid_prev2_node); - // gamma is a runtime [1] scalar tensor -> broadcast-multiply. - ggml_tensor* scaled = ggml_mul(ctx, diff, gamma_scalar); - ggml_tensor* aligned = ggml_add(ctx, resid_prev2_node, scaled); - return ggml_add(ctx, x_before, aligned); - } - // GPU-side single-residual reuse (MagCache/TaylorSeer feature reuse): - // inject the last captured residual straight from device memory, avoiding - // the ~50MB host reconstruct copy + H2D upload the inject_feature path pays. - if (gpu_reuse_available()) { - return ggml_add(ctx, x_before, resid_prev1_node); - } - return ggml_add(ctx, x_before, inject_feature); - } - bool gpu_reconstruct_available() const { - return resid_prev1_node != nullptr && resid_prev2_node != nullptr && - gamma_scalar != nullptr; - } - // Single-residual on-GPU reuse: resid_prev1 present, no ring/gamma (MagCache). - bool gpu_reuse_available() const { - return resid_prev1_node != nullptr && resid_prev2_node == nullptr && - gamma_scalar == nullptr; - } - void build_feature(ggml_context* ctx, ggml_tensor* x_after) { - if (before_node != nullptr) { - feature_node = ggml_sub(ctx, x_after, before_node); - } - } - // Build the GPU-side DiCache decision scalars from the probe state. Called by - // the runner after the model records before_node/probe_node on a Probe pass, - // when gpu_metric is set and the persistent prev_* nodes are available. - // delta_y = sum|probe - prev_probe| / sum|prev_probe| - // delta_x = sum|before - prev_input| / sum|prev_input| (delta_minus only) - // gamma = clamp( sum|(probe-before) - probe_prev2| / sum|probe_prev1 - probe_prev2|, 1, 1.5) - // (mean's 1/N cancels in each ratio, so sum is used directly.) - void build_probe_metrics(ggml_context* ctx) { - if (!gpu_metric || probe_node == nullptr || before_node == nullptr) { - return; - } - auto rel = [&](ggml_tensor* cur, ggml_tensor* ref) -> ggml_tensor* { - ggml_tensor* num = ggml_sum(ctx, ggml_abs(ctx, ggml_sub(ctx, cur, ref))); - ggml_tensor* den = ggml_sum(ctx, ggml_abs(ctx, ref)); - return ggml_div(ctx, num, den); - }; - if (prev_probe_node != nullptr) { - delta_y_node = rel(probe_node, prev_probe_node); - } - if (prev_input_node != nullptr) { - delta_x_node = rel(before_node, prev_input_node); - } - if (probe_prev1_node != nullptr && probe_prev2_node != nullptr) { - ggml_tensor* cur_resid = ggml_sub(ctx, probe_node, before_node); - ggml_tensor* num = ggml_sum(ctx, ggml_abs(ctx, ggml_sub(ctx, cur_resid, probe_prev2_node))); - ggml_tensor* den = ggml_sum(ctx, ggml_abs(ctx, ggml_sub(ctx, probe_prev1_node, probe_prev2_node))); - gamma_node = ggml_div(ctx, num, den); - } - } - - - // ---- In-loop sugar (per-block region boundaries) ---- - // Call before block i runs: snapshot the region-input anchor iff i is the - // region's first block and we are capturing/probing (not injecting). - void begin_region(int i, ggml_tensor* x) { - if (!inject_mode() && at_region_start(i)) { - snapshot_before(x); - } - } - // Call at the top of iteration i under Inject: if i is the region start, - // returns region_before + inject_feature (caller assigns it and jumps to - // inject_resume_index). Returns nullptr when block i should just run. - ggml_tensor* step_inject_region(ggml_context* ctx, int i, ggml_tensor* x) const { - if (inject_mode() && at_region_start(i)) { - return add_injected(ctx, x); - } - return nullptr; - } - // GPU DiCache: during a Capture (full compute) step, also snapshot the - // hidden state after probe_depth blocks so the runner can refresh the - // persistent prev_probe / probe-residual used by the NEXT step's decision. - // No effect unless capturing with gpu_metric and probe_depth>0. - void record_capture_probe_state(ggml_context* ctx, int i, ggml_tensor* x) { - if (capture_mode() && gpu_metric && probe_depth > 0 && - (i - region_start) + 1 == probe_depth) { - probe_node = x; - // Probe residual (probe_state - block_input) for the ring; before_node - // is the whole-stack input, snapshotted at region start. - if (before_node != nullptr) { - probe_resid_node = ggml_sub(ctx, x, before_node); - } - } - } - // Call after block i runs: on Capture, build the region residual once i is - // the region's last block. - void end_region(ggml_context* ctx, int i, int n, ggml_tensor* x) { - if (mode == Mode::Capture && before_node != nullptr && i + 1 == resolved_region_end(n)) { - build_feature(ctx, x); - } - } - - // True when a Probe should stop right after this 0-based block index - // (counted from region_start). - bool stop_after_block(int i) const { - return mode == Mode::Probe && (i - region_start) + 1 >= probe_depth; - } - void on_probe(ggml_tensor* x_probe) { probe_node = x_probe; } -}; - } // namespace sd diff --git a/src/core/optimization/cache/compile/cache_graph_lowering.cpp b/src/core/optimization/cache/compile/cache_graph_lowering.cpp index 8a5669e2..84f51bd5 100644 --- a/src/core/optimization/cache/compile/cache_graph_lowering.cpp +++ b/src/core/optimization/cache/compile/cache_graph_lowering.cpp @@ -148,20 +148,6 @@ class ActionInterpreter { std::unordered_map> temps_; }; -// The (start,end) region the block-stack segment covers, for capture/inject. -void seam_region(const GraphVariantPlan& v, int* start, int* end, int* probe_depth) { - *start = 0; - *end = -1; - *probe_depth = 0; - for (const auto& seg : v.segments) { - if (seg.execution == SegmentExecutionMode::PROBE) { - *probe_depth = seg.probe_depth; - } - *start = seg.region_start; - *end = seg.region_end; - } -} - } // namespace sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, @@ -179,7 +165,7 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, const int slot0 = program.slots.empty() ? 0 : program.slots.front().id; const bool device_slot = state.has_device_store() && !program.slots.empty() && program.slots.front().device_backed && - hooks.capture_to_slot && hooks.inject_from_slot; + hooks.substep_capture && hooks.substep_inject_slot; sd::Tensor out; while (auto plan_opt = policy.next_substep()) { @@ -187,8 +173,6 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, SubstepResult result; sd::Tensor y; - const int region_end = plan.blocks.end; // -1 => to end of stack - // ---- Output-granularity substep (EasyCache/UCache/Condition): the reuse / // capture works on the whole denoiser output via the declarative operators // (LOAD/STORE/DIFFERENCE/BLEND over the state-manager slot), not the block @@ -246,11 +230,6 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, // Mirrors execute_declarative_feature's host path, driven by next_substep. ---- if (plan.op.kind == SubstepOpKind::FeatureReuse || plan.op.kind == SubstepOpKind::FeatureCompute) { - int fr_start = 0, fr_end = -1, fr_probe = 0; - const GraphVariantPlan* fr_variant = program.find_kind(GraphVariantKind::FULL); - if (fr_variant != nullptr) { - seam_region(*fr_variant, &fr_start, &fr_end, &fr_probe); - } ActionInterpreter interp(state, operators, condition_key); interp.bind_ambient(kAmbientInput, hooks.input); interp.set_step_coeffs(&plan.op.coeffs); @@ -264,13 +243,16 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, if (!interp.run(seg.before)) { ok = false; break; } } const sd::Tensor* feat = ok ? interp.ambient(kAmbientInjectFeature) : nullptr; - if (feat != nullptr && !feat->empty() && hooks.inject) { - y = hooks.inject(*feat, fr_start, fr_end); + if (feat != nullptr && !feat->empty() && hooks.substep_inject_host) { + // Tap-driven host inject (Wan/SD3, no CacheGraphScope). + y = hooks.substep_inject_host(*feat); } // Ring not ready / inject failed -> capture below. } - if (y.empty() && hooks.capture) { - sd::DiffusionCacheResult res = hooks.capture(fr_start, fr_end); + if (y.empty() && hooks.substep_capture_host) { + // Tap-driven host capture (Wan/SD3, no CacheGraphScope). Feature methods + // are always whole-stack, so the no-arg host variant is exact. + sd::DiffusionCacheResult res = hooks.substep_capture_host(); y = std::move(res.output); if (!y.empty() && !res.feature.empty()) { CacheObservation obs; @@ -309,7 +291,7 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, ? plan.blocks.end : std::max(1, plan.blocks.begin); if (hooks.substep_probe) { - // Tap-driven probe (ED_CACHE_SUBSTEP): delta_y/gamma computed + // Tap-driven probe: delta_y/gamma computed // on-device from taps + persistent operands, no CacheGraphScope. sd::DiffusionCacheResult pr = hooks.substep_probe(depth); result.indicators["delta_y"] = pr.delta_y; @@ -322,12 +304,6 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, // host (reusing decide_after_probe). No CacheGraphScope. sd::DiffusionCacheResult pr = hooks.substep_probe_host(depth); policy.observe_substep_probe_host(&pr.before, &pr.probe, step, condition_key); - } else if (hooks.probe) { - sd::DiffusionCacheResult pr = hooks.probe(depth); - result.indicators["delta_y"] = pr.delta_y; - result.indicators["delta_x"] = pr.delta_x; - result.indicators["gamma"] = pr.gamma; - policy.observe_substep(result); } continue; // probe never produces the step output } @@ -337,12 +313,9 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, if (plan.op.kind == SubstepOpKind::Extrapolate) { if (!plan.op.coeffs.empty()) { const float gamma = plan.op.coeffs.front(); - // Prefer the tap-driven device inject (no CacheGraphScope); fall back - // to the legacy inject_gpu only if the tap path isn't wired. + // Tap-driven device inject (no CacheGraphScope). if (hooks.substep_inject_gpu) { y = hooks.substep_inject_gpu(gamma, 0, -1); - } else if (hooks.inject_gpu) { - y = hooks.inject_gpu(gamma, 0, -1); } } if (y.empty()) { @@ -356,6 +329,48 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, continue; } + // ---- DiCache seed compute: full forward that refreshes the cross-step + // residual/probe rings. Device path (Qwen/Flux) writes them on-device via the + // tap-driven substep_capture_probe; host path (Wan/SD3) reads the feature back + // for the policy's host ring. Replaces the legacy hooks.capture path (which ran + // run_cache_pass + CacheGraphScope). Produces the step output. ---- + if (plan.op.kind == SubstepOpKind::CaptureProbeSeed) { + if (hooks.substep_capture_probe) { + y = hooks.substep_capture_probe(); + if (!y.empty()) { + CacheObservation obs; + obs.kind = CacheObservation::Kind::Feature; + obs.step = step; + obs.condition_key = condition_key; + obs.branch = branch; + obs.input = hooks.input; + obs.feature_on_device = true; // DiCache observe() is a no-op on the device path + policy.observe(obs); + } + } else if (hooks.substep_capture_host) { + sd::DiffusionCacheResult res = hooks.substep_capture_host(); + y = std::move(res.output); + if (!y.empty() && !res.feature.empty()) { + CacheObservation obs; + obs.kind = CacheObservation::Kind::Feature; + obs.step = step; + obs.condition_key = condition_key; + obs.branch = branch; + obs.input = hooks.input; + obs.feature = &res.feature; + policy.observe(obs); + } + } + if (y.empty()) { + y = hooks.full(); + } + policy.observe_substep(result); + if (plan.produces_output) { + out = std::move(y); + } + continue; + } + const bool is_reuse = plan.op.kind == SubstepOpKind::ApplyResidual; // Seam region semantics: SubstepPlan.blocks is the interval to COMPUTE. The // inject/capture region is the block interval the residual spans. For the @@ -368,25 +383,21 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, // capture below. CacheSlotHandle h = state.read(condition_key, slot0); if (h.valid && h.buffer != nullptr) { - // Prefer the tap-driven device inject (no CacheGraphScope); fall back - // to the legacy inject_from_slot only if the tap path isn't wired. - y = hooks.substep_inject_slot ? hooks.substep_inject_slot(h.buffer, 0, -1) - : hooks.inject_from_slot(h.buffer, 0, -1); + // Tap-driven device inject (no CacheGraphScope). + y = hooks.substep_inject_slot(h.buffer, 0, -1); } - } else if (is_reuse && (hooks.substep_inject_host || hooks.inject)) { - // Host reuse (no device slot — Wan): the policy reconstructs the residual - // on host; inject it as x_before + feature over the whole stack. Empty - // reconstruct (history not ready) falls through to a capture below. + } else if (is_reuse && hooks.substep_inject_host) { + // Host reuse (no device slot — Wan/SD3): the policy reconstructs the + // residual on host; inject it as x_before + feature over the whole stack. + // Empty reconstruct (history not ready) falls through to a capture below. CacheReconstructContext rc; rc.step = step; rc.condition_key = condition_key; rc.input = hooks.input; sd::Tensor feature = policy.reconstruct(rc); if (!feature.empty()) { - // Prefer the tap-driven inject (no CacheGraphScope); fall back to the - // legacy compute_inject only if the tap path isn't wired. - y = hooks.substep_inject_host ? hooks.substep_inject_host(feature) - : hooks.inject(feature, 0, -1); + // Tap-driven host inject (no CacheGraphScope). + y = hooks.substep_inject_host(feature); } } @@ -395,7 +406,7 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, // into the device slot when this plan writes it. const bool wants_capture = !plan.writes.empty(); if (device_slot && wants_capture && hooks.substep_capture) { - // Tap-driven capture (ED_CACHE_SUBSTEP): the runner requests + // Tap-driven capture: the runner requests // ModelIn/ModelOut taps and weaves the residual — no CacheGraphScope. auto alloc_slot = [&](const std::vector& residual_shape) -> void* { return state.alloc_device_entry(condition_key, slot0, residual_shape); @@ -411,23 +422,8 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, obs.feature_on_device = true; policy.observe(obs); } - } else if (device_slot && wants_capture && hooks.capture_to_slot) { - auto alloc_slot = [&](const std::vector& residual_shape) -> void* { - return state.alloc_device_entry(condition_key, slot0, residual_shape); - }; - y = hooks.capture_to_slot(alloc_slot, 0, region_end); - if (!y.empty()) { - CacheObservation obs; - obs.kind = CacheObservation::Kind::Feature; - obs.step = step; - obs.condition_key = condition_key; - obs.branch = branch; - obs.input = hooks.input; - obs.feature_on_device = true; - policy.observe(obs); - } } else if (wants_capture && hooks.substep_capture_host) { - // Tap-driven host capture (ED_CACHE_SUBSTEP, no device slot — Wan): + // Tap-driven host capture (no device slot — Wan/SD3): // the runner weaves (ModelOut - ModelIn) and reads it back to host. // No CacheGraphScope. The policy stores the host residual for reuse. sd::DiffusionCacheResult res = hooks.substep_capture_host(); @@ -442,22 +438,6 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, obs.feature = &res.feature; policy.observe(obs); } - } else if (wants_capture && hooks.capture) { - // Host/GPU-ring capture (DiCache): compute_capture runs the full - // forward AND seeds the cross-step residual/probe ring via its - // handoff. Must be used instead of full() so the ring updates. - sd::DiffusionCacheResult res = hooks.capture(0, region_end); - y = std::move(res.output); - if (!y.empty() && !res.feature.empty()) { - CacheObservation obs; - obs.kind = CacheObservation::Kind::Feature; - obs.step = step; - obs.condition_key = condition_key; - obs.branch = branch; - obs.input = hooks.input; - obs.feature = &res.feature; - policy.observe(obs); - } } if (y.empty()) { // No capture wired / capture failed / plain compute substep. diff --git a/src/core/optimization/cache/compile/cache_graph_lowering.hpp b/src/core/optimization/cache/compile/cache_graph_lowering.hpp index 175b7fb5..3a25b151 100644 --- a/src/core/optimization/cache/compile/cache_graph_lowering.hpp +++ b/src/core/optimization/cache/compile/cache_graph_lowering.hpp @@ -14,50 +14,32 @@ namespace edgedit { namespace cache { // Model-agnostic compute callbacks the pipeline wires to its exact runner calls. -// Unchanged in spirit from the pre-redesign CacheRunnerHooks: the lowering -// invokes whichever subset the chosen variant needs. -// full - normal full transformer forward (also every fallback path). -// input - the block-stack input latent (Output-level diff + SenCache). -// capture/inject/probe - the Feature/Probe seam passes. -// feature_supported is false when the model can't cut its block stack this run. +// The runner entry points the lowering invokes to run one substep. All cache +// execution is tap-driven (TapRegistry, no CacheGraphScope); the lowering calls +// whichever subset the chosen SubstepPlan needs. +// full - normal full transformer forward (also every fallback path). +// input - the block-stack input latent (Output-level diff + SenCache). +// substep_* - the tap-driven capture/probe/inject passes (device or host). struct CacheRunnerHooks { const sd::Tensor* input = nullptr; std::function()> full; - bool feature_supported = false; - std::function capture; // (region_start, region_end) - std::function(const sd::Tensor&, int, int)> inject; // (feature, start, end) - std::function probe; // (probe_depth) - // GPU DiCache reuse: reconstruct + inject on-device from the persistent - // residual ring using a host-clamped gamma. (gamma, region_start, region_end). - // Empty result => insufficient history; lowering falls back to a full step. - std::function(float, int, int)> inject_gpu; - // ---- Declarative device-slot seam (B2). When the CacheStateManager backs a - // Feature slot with a device tensor, the lowering drives store/reuse through - // these (the on-GPU feature-reuse path for MagCache/TaylorSeer). The void* is - // the slot's opaque ggml_tensor* (from CacheSlotHandle::buffer). - // capture_to_slot: full compute; AFTER the block-stack residual shape is - // known, calls alloc_slot(residual_shape) to obtain the device slot tensor - // (the StateManager allocates it at that shape), copies the residual into - // it on-device, and returns the model output (empty => fail). Passing the - // allocator (not a pre-sized slot) is required because the residual shape - // is the packed block-stack seq shape, unknown to the lowering pre-capture. - // inject_from_slot: reuse step injecting x_before + slot residual on-device; - // returns output (empty => slot not ready -> caller falls back to capture). - std::function( - const std::function&)>&, int, int)> capture_to_slot; - std::function(void*, int, int)> inject_from_slot; // (slot, start, end) - // Substep-path (ED_CACHE_SUBSTEP) tap-driven capture: same contract as - // capture_to_slot (alloc_slot returns the device slot for the residual shape) - // but the runner drives it through the TapRegistry, not CacheGraphScope. Null - // on the legacy path. + // Substep-path tap-driven capture (MagCache device slot): same contract as a + // device-slot store — alloc_slot returns the device slot for the residual + // shape; the runner drives it through the TapRegistry. std::function( const std::function&)>&)> substep_capture; - // Substep-path (ED_CACHE_SUBSTEP) tap-driven DiCache probe: runs the shallow + // Substep-path tap-driven DiCache device seed capture: full forward whose + // post-readback refreshes the persistent DiCacheGpuState residual/probe rings + // device-to-device from the tap-woven nodes (no device slot — the ring is + // runner-owned, not a StateManager slot). Returns the step output. + // Null on the host path. + std::function()> substep_capture_probe; + + // Substep-path tap-driven DiCache probe: runs the shallow // prefix and returns delta_y/delta_x/gamma computed on-device from taps + - // persistent operands, via the TapRegistry (no CacheGraphScope). (probe_depth). - // Null on the legacy path. + // persistent operands, via the TapRegistry. (probe_depth). Null on the host path. std::function substep_probe; // Host-path substep variants (models with no device slot — Wan). Tap-driven like @@ -75,10 +57,9 @@ struct CacheRunnerHooks { std::function(float, int, int)> substep_inject_gpu; // (gamma, start, end) }; -// Executes one chosen GraphVariantPlan against the runner hooks, using the -// policy for host-side reconstruction and observation. This is the seam between -// the declarative program and the current CacheGraphScope-based runner passes -// (Option A lowering — see the plan's "Key architectural decision"). +// Drives the policy's substep loop against the runner hooks, using the policy for +// host-side reconstruction and observation. This is the seam between the +// declarative program and the tap-driven runner passes. class CacheGraphLowering { public: // The substep loop: drive the policy's next_substep() for one CFG branch, diff --git a/src/core/optimization/cache/ir/substep_plan.hpp b/src/core/optimization/cache/ir/substep_plan.hpp index 9ca09d08..b75ae2f0 100644 --- a/src/core/optimization/cache/ir/substep_plan.hpp +++ b/src/core/optimization/cache/ir/substep_plan.hpp @@ -39,6 +39,8 @@ enum class SubstepOpKind { OutputCompute, // full forward, then STORE (output - input) diff to slot [Output methods] FeatureReuse, // host feature-ring reuse: PREDICT/REUSE before-actions -> inject [TaylorSeer/SenCache] FeatureCompute, // full forward via hooks.capture, then STORE+ROTATE ring [TaylorSeer/SenCache] + CaptureProbeSeed,// DiCache seed: full forward + tap-driven writeback of the cross-step + // residual/probe rings (device) or host feature readback [DiCache] }; // Which declarative variant the Output-method substep drives (whole-output diff), diff --git a/src/core/optimization/cache/model/tap_registry.hpp b/src/core/optimization/cache/model/tap_registry.hpp index ee3841a6..94137f51 100644 --- a/src/core/optimization/cache/model/tap_registry.hpp +++ b/src/core/optimization/cache/model/tap_registry.hpp @@ -68,6 +68,8 @@ class TapRegistry { indicators_.clear(); stop_after_block_ = -1; capture_residual_ = false; + capture_writeback_ = false; + writeback_probe_depth_ = 0; inject_active_ = false; inject_kind_ = InjectKind::None; inject_input_ = nullptr; @@ -99,6 +101,21 @@ class TapRegistry { void set_capture_residual(bool on) { capture_residual_ = on; } bool capture_residual() const { return capture_residual_; } + // Capture-writeback (DiCache device seed): on top of the ModelOut-ModelIn + // residual (set_capture_residual), also weave the probe residual + // (BlockOut[probe_depth-1] - ModelIn) as a named node "ed_cache_probe_resid" so + // the pass's post-readback can d2d it (plus the raw probe/before taps) into the + // runner's persistent DiCache ring. Symmetric to set_probe_metrics but for the + // capture direction. Only probe_depth is needed here: the destination ring + // tensors are allocated by the runner AFTER forward (their shape is the packed + // block-stack seq shape, unknown at set time), so they can't be passed in now. + void set_capture_writeback(int probe_depth) { + capture_writeback_ = true; + writeback_probe_depth_ = probe_depth; + } + bool capture_writeback() const { return capture_writeback_; } + int writeback_probe_depth() const { return writeback_probe_depth_; } + // ---- Tap-driven inject (reuse). Replaces cache_scope->step_inject_region in the // model forward: when active, at block `inject_start` the forward replaces the // stream with (x_before + ) and jumps to `inject_resume` (skipping the @@ -176,6 +193,8 @@ class TapRegistry { std::vector indicators_; int stop_after_block_ = -1; bool capture_residual_ = false; + bool capture_writeback_ = false; + int writeback_probe_depth_ = 0; bool inject_active_ = false; InjectKind inject_kind_ = InjectKind::None; ggml_tensor* inject_input_ = nullptr; // HostFeature diff --git a/src/core/optimization/cache/policy/cache_policy.hpp b/src/core/optimization/cache/policy/cache_policy.hpp index 67fa587a..2fd57a03 100644 --- a/src/core/optimization/cache/policy/cache_policy.hpp +++ b/src/core/optimization/cache/policy/cache_policy.hpp @@ -137,19 +137,12 @@ class ICachePolicy { return decide(step, {}); } - // ---- Substep interface (ED_CACHE_SUBSTEP path) -------------------------- - // The clean-slate model: a denoise step is a sequence of substeps, each a + // ---- Substep interface -------------------------------------------------- + // The execution model: a denoise step is a sequence of substeps, each a // decide->build->run unit. next_substep() is pure host scalar logic (no ggml, // no device), so it is unit-testable off-hardware. Most methods yield one - // substep; DiCache yields two (probe, then continuation). Default: opt out, so - // the engine keeps using decide()/decide_after_probe() for unmigrated methods. - virtual bool supports_substep() const { return false; } - - // The engine reports whether an on-device metric path is available this run (a - // device store was wired by the runner). Methods whose substep path needs - // on-device reductions (DiCache) gate supports_substep() on this; host-only - // methods ignore it. Called once after compile(). Default no-op. - virtual void set_substep_device_available(bool /*available*/) {} + // substep; DiCache yields two (probe, then continuation). This is the sole + // cache execution path — run_branch always drives the substep loop. // The engine hands the block-stack input latent (hooks.input) before the substep // loop, so Output methods can run their decide() metric. Default no-op. diff --git a/src/core/optimization/cache/policy/policies/condition_policy.cpp b/src/core/optimization/cache/policy/policies/condition_policy.cpp index 006c12c0..5bc2f0aa 100644 --- a/src/core/optimization/cache/policy/policies/condition_policy.cpp +++ b/src/core/optimization/cache/policy/policies/condition_policy.cpp @@ -64,8 +64,6 @@ class ConditionCachePolicy final : public ICachePolicy { // via the residual-diff slot; the dead taylor_states_ bookkeeping in // observe() is harmless. declarative_ = true; - const char* substep_env = std::getenv("ED_CACHE_SUBSTEP"); - substep_env_on_ = substep_env != nullptr && substep_env[0] != '\0' && substep_env[0] != '0'; return detail::make_output_diff_program(method_label_for_mode(mode_), seg); } @@ -199,9 +197,8 @@ class ConditionCachePolicy final : public ICachePolicy { return out; } - // ---- Substep interface (ED_CACHE_SUBSTEP). Output method (DBCache/CacheDiT): - // one substep, reuse (declarative diff) or compute; decision reuses decide(). ---- - bool supports_substep() const override { return substep_env_on_; } + // ---- Substep interface. Output method (DBCache/CacheDiT): one substep, + // reuse (declarative diff) or compute; decision reuses decide(). ---- void set_substep_input(const sd::Tensor* input) override { substep_input_ = input; } void begin_substeps(const StepContext& step, const void* condition_key) override { begin_step(step); @@ -374,7 +371,6 @@ class ConditionCachePolicy final : public ICachePolicy { const void* anchor_condition_ = nullptr; std::unordered_map cache_entries_; std::unordered_map taylor_states_; - bool substep_env_on_ = false; bool substep_done_ = false; const void* substep_key_ = nullptr; const sd::Tensor* substep_input_ = nullptr; diff --git a/src/core/optimization/cache/policy/policies/dicache_policy.cpp b/src/core/optimization/cache/policy/policies/dicache_policy.cpp index 03ec922e..cf0bff82 100644 --- a/src/core/optimization/cache/policy/policies/dicache_policy.cpp +++ b/src/core/optimization/cache/policy/policies/dicache_policy.cpp @@ -1,6 +1,5 @@ #include #include -#include #include #include @@ -49,36 +48,23 @@ class DiCachePolicy final : public ICachePolicy { total_steps_skipped_ = 0; const int seg = topo.block_stack() ? topo.block_stack()->id : 1; - // Declarative Probe program: residual ring + PROBE/REUSE/FULL variants. On - // the host path (ED_DICACHE_GPU=0, so the pipeline leaves inject_gpu unset) - // the lowering runs the probe, calls decide_after_probe (which supplies the - // gamma blend weights), and injects the declarative blend of the ring. The - // default on-GPU path leaves inject_gpu set, so it takes the legacy seam - // control flow (which this same program also drives). - // Substep path opt-in: gated by ED_CACHE_SUBSTEP AND a viable on-device - // metric path (a device store wired by the runner). Wan has no device store, - // so its host DiCache stays on the legacy path. The engine calls - // set_substep_gpu_path() after init with the device-store signal; here we - // only record the env opt-in, ANDed in supports_substep(). - const char* substep_env = std::getenv("ED_CACHE_SUBSTEP"); - substep_env_on_ = substep_env != nullptr && substep_env[0] != '\0' && substep_env[0] != '0'; - substep_gpu_path_ = false; // set true by set_substep_gpu_path() when a store exists + // Declarative Probe program: residual ring + PROBE/REUSE/FULL variants. The + // on-device metric path (Qwen/Flux, device store) surfaces delta_y/gamma from + // the probe seam as scalars; the host path (Wan, no store, or ED_DICACHE_GPU=0) + // runs the probe, calls decide_after_probe (which supplies the gamma blend + // weights), and injects the declarative blend of the residual ring. return detail::make_dicache_program("DiCache", seg, std::max(1, config_.probe_depth)); } void begin_step(const StepContext& step) override { current_step_index_ = step.step_index; } - // ---- Substep interface (ED_CACHE_SUBSTEP). DiCache is the two-yield case that - // proves the uniform loop: substep 1 = probe (measures delta_y/gamma on-device, - // produces no output), then observe_substep folds the scalars via the same + // ---- Substep interface. DiCache is the two-yield case that proves the + // uniform loop: substep 1 = probe (measures delta_y/gamma on-device, produces + // no output), then observe_substep folds the scalars via the same // decide_after_probe math; substep 2 = reuse (Extrapolate with gamma) or a - // capturing full compute. Both paths are migrated: the on-device metric path + // capturing full compute. Both paths are supported: the on-device metric path // (Qwen/Flux, device store) and the tap-driven host path (Wan, no store, uses - // observe_substep_probe_host + host reconstruct). Gate only on the env opt-in. - bool supports_substep() const override { return substep_env_on_; } - - void set_substep_device_available(bool available) override { substep_gpu_path_ = available; } - + // observe_substep_probe_host + host reconstruct). void begin_substeps(const StepContext& step, const void* condition_key) override { current_step_index_ = step.step_index; substep_key_ = condition_key; @@ -98,6 +84,10 @@ class DiCachePolicy final : public ICachePolicy { p.blocks = BlockRange{0, -1}; p.produces_output = true; p.writes = {0}; // capture residual into the ring + // Seed the cross-step rings on this full compute too (retention-prefix + // and final steps): same tap-driven writeback as the probe-eligible + // seed, so no step falls to the legacy hooks.capture path. + p.op = SubstepOp{SubstepOpKind::CaptureProbeSeed, {}, {}}; return p; } case SubstepPhase::Probe: { @@ -131,6 +121,10 @@ class DiCachePolicy final : public ICachePolicy { } else { p.blocks = BlockRange{0, -1}; p.writes = {0}; + // Seed compute: refresh the cross-step rings via the tap-driven + // capture. Device (Qwen/Flux) writes the DiCacheGpuState ring; + // host (Wan/SD3) falls to substep_capture_host in the lowering. + p.op = SubstepOp{SubstepOpKind::CaptureProbeSeed, {}, {}}; } return p; } @@ -443,10 +437,8 @@ class DiCachePolicy final : public ICachePolicy { int total_steps_skipped_ = 0; std::unordered_map states_; - // Substep state (ED_CACHE_SUBSTEP GPU path). + // Substep state. enum class SubstepPhase { Probe, Continue, SingleFull, Done }; - bool substep_env_on_ = false; - bool substep_gpu_path_ = false; SubstepPhase substep_phase_ = SubstepPhase::Done; const void* substep_key_ = nullptr; bool substep_reuse_ = false; diff --git a/src/core/optimization/cache/policy/policies/easycache_policy.cpp b/src/core/optimization/cache/policy/policies/easycache_policy.cpp index 0cb90a7e..104b0ac1 100644 --- a/src/core/optimization/cache/policy/policies/easycache_policy.cpp +++ b/src/core/optimization/cache/policy/policies/easycache_policy.cpp @@ -74,8 +74,6 @@ class EasyCachePolicy final : public ICachePolicy { window_.configure(inf.sigmas ? *inf.sigmas : std::vector{}, config_.start_percent, config_.end_percent); reset(); const int seg = topo.block_stack() ? topo.block_stack()->id : 1; - const char* substep_env = std::getenv("ED_CACHE_SUBSTEP"); - substep_env_on_ = substep_env != nullptr && substep_env[0] != '\0' && substep_env[0] != '0'; // Declarative program: the residual slot + difference/blend operators are // driven by the lowering. This policy now only owns its scalar decision // state (prev input/output metrics); the cached diff tensor lives in the @@ -143,10 +141,8 @@ class EasyCachePolicy final : public ICachePolicy { return {}; } - // ---- Substep interface (ED_CACHE_SUBSTEP). Output method: one substep per - // step, reuse (whole-output diff) or compute; decision reuses decide(). ---- - bool supports_substep() const override { return substep_env_on_; } - + // ---- Substep interface. Output method: one substep per step, reuse + // (whole-output diff) or compute; decision reuses decide(). ---- void begin_substeps(const StepContext& step, const void* condition_key) override { begin_step(step); substep_done_ = false; @@ -285,7 +281,6 @@ class EasyCachePolicy final : public ICachePolicy { bool has_last_input_change_ = false; int total_steps_skipped_ = 0; int current_step_index_ = -1; - bool substep_env_on_ = false; bool substep_done_ = false; const void* substep_key_ = nullptr; const sd::Tensor* substep_input_ = nullptr; diff --git a/src/core/optimization/cache/policy/policies/magcache_policy.cpp b/src/core/optimization/cache/policy/policies/magcache_policy.cpp index 1d566733..d6317673 100644 --- a/src/core/optimization/cache/policy/policies/magcache_policy.cpp +++ b/src/core/optimization/cache/policy/policies/magcache_policy.cpp @@ -217,9 +217,6 @@ class MagCachePolicy final : public ICachePolicy { // retains only its scalar decision state (accumulated ratio / has_residual) // and, on a calibrate run, the previous-step feature. (void)seg; - const char* substep_env = std::getenv("ED_CACHE_SUBSTEP"); - substep_env_on_ = substep_env != nullptr && substep_env[0] != '\0' && substep_env[0] != '0'; - substep_device_available_ = false; // set by set_substep_device_available() return detail::make_feature_reuse_program("MagCache", seg, /*device_backed=*/true); } @@ -237,13 +234,8 @@ class MagCachePolicy final : public ICachePolicy { // ---- Substep interface. Same decision math as decide(); reuse => a single // zero-block ApplyResidual substep, compute => a full-stack substep that // captures the residual slot. The GPU feature-reuse path (device slot) is - // driven by the adapter exactly as before; the host substep path (no store — - // Wan) reconstructs the residual on host and injects via hooks.inject. Both are - // supported, so gate only on the env opt-in. - bool supports_substep() const override { return substep_env_on_; } - - void set_substep_device_available(bool available) override { substep_device_available_ = available; } - + // driven by the adapter; the host substep path (no store — Wan) reconstructs + // the residual on host and injects via hooks.inject. Both are supported. void begin_substeps(const StepContext& step, const void* condition_key) override { current_step_index_ = step.step_index; substep_done_ = false; @@ -495,8 +487,6 @@ class MagCachePolicy final : public ICachePolicy { std::unordered_map states_; bool substep_done_ = false; const void* substep_branch_key_ = nullptr; - bool substep_env_on_ = false; - bool substep_device_available_ = false; Branch& branch_for(const void* cond) { return states_[cond]; } }; diff --git a/src/core/optimization/cache/policy/policies/sencache_policy.cpp b/src/core/optimization/cache/policy/policies/sencache_policy.cpp index 6118126e..82f83ef7 100644 --- a/src/core/optimization/cache/policy/policies/sencache_policy.cpp +++ b/src/core/optimization/cache/policy/policies/sencache_policy.cpp @@ -171,9 +171,6 @@ class SenCachePolicy final : public ICachePolicy { ? detail::make_feature_capture_program("SenCache", seg) : detail::make_feature_reuse_program("SenCache", seg, /*device_backed=*/true); - const char* substep_env = std::getenv("ED_CACHE_SUBSTEP"); - substep_env_on_ = substep_env != nullptr && substep_env[0] != '\0' && substep_env[0] != '0'; - if (calibrating_) { initialized_ = true; LOG_INFO("SenCache calibration: profiling %d steps -> %s", @@ -289,10 +286,9 @@ class SenCachePolicy final : public ICachePolicy { return {}; } - // ---- Substep interface (ED_CACHE_SUBSTEP). Feature method (single-residual - // reuse): one substep, reuse (LOAD slot -> inject) or compute (capture + STORE); + // ---- Substep interface. Feature method (single-residual reuse): one + // substep, reuse (LOAD slot -> inject) or compute (capture + STORE); // decision reuses decide(). ---- - bool supports_substep() const override { return substep_env_on_; } void set_substep_input(const sd::Tensor* input) override { substep_input_ = input; } void begin_substeps(const StepContext& step, const void* condition_key) override { begin_step(step); @@ -505,7 +501,6 @@ class SenCachePolicy final : public ICachePolicy { std::vector sigmas_; SenCacheProfile profile_; std::unordered_map states_; - bool substep_env_on_ = false; bool substep_done_ = false; const void* substep_key_ = nullptr; const sd::Tensor* substep_input_ = nullptr; diff --git a/src/core/optimization/cache/policy/policies/taylorseer_policy.cpp b/src/core/optimization/cache/policy/policies/taylorseer_policy.cpp index 8ae1afe8..8d28cb97 100644 --- a/src/core/optimization/cache/policy/policies/taylorseer_policy.cpp +++ b/src/core/optimization/cache/policy/policies/taylorseer_policy.cpp @@ -53,8 +53,6 @@ class TaylorSeerFeaturePolicy final : public ICachePolicy { depth_ = order_ + 2; // must match make_taylor_history_program(depth = order + 2) reset(); const int seg = topo.block_stack() ? topo.block_stack()->id : 1; - const char* substep_env = std::getenv("ED_CACHE_SUBSTEP"); - substep_env_on_ = substep_env != nullptr && substep_env[0] != '\0' && substep_env[0] != '0'; return detail::make_taylor_history_program("TaylorSeer", seg, order_); } @@ -106,10 +104,9 @@ class TaylorSeerFeaturePolicy final : public ICachePolicy { // serves the injected feature from the ring. sd::Tensor reconstruct(const CacheReconstructContext&) override { return {}; } - // ---- Substep interface (ED_CACHE_SUBSTEP). Feature method: one substep, reuse - // (host feature-ring PREDICT blend) or compute (capture+STORE+ROTATE); decision + // ---- Substep interface. Feature method: one substep, reuse (host + // feature-ring PREDICT blend) or compute (capture+STORE+ROTATE); decision // and blend coeffs reuse decide(). ---- - bool supports_substep() const override { return substep_env_on_; } void begin_substeps(const StepContext& step, const void* condition_key) override { begin_step(step); substep_done_ = false; @@ -274,7 +271,6 @@ class TaylorSeerFeaturePolicy final : public ICachePolicy { int order_ = 1; int depth_ = 3; std::unordered_map states_; - bool substep_env_on_ = false; bool substep_done_ = false; const void* substep_key_ = nullptr; StepContext substep_step_; diff --git a/src/core/optimization/cache/policy/policies/ucache_policy.cpp b/src/core/optimization/cache/policy/policies/ucache_policy.cpp index 24fff12c..86aad112 100644 --- a/src/core/optimization/cache/policy/policies/ucache_policy.cpp +++ b/src/core/optimization/cache/policy/policies/ucache_policy.cpp @@ -41,8 +41,6 @@ class UCachePolicy final : public ICachePolicy { window_.configure(inf.sigmas ? *inf.sigmas : std::vector{}, config_.start_percent, config_.end_percent); reset(); const int seg = topo.block_stack() ? topo.block_stack()->id : 1; - const char* substep_env = std::getenv("ED_CACHE_SUBSTEP"); - substep_env_on_ = substep_env != nullptr && substep_env[0] != '\0' && substep_env[0] != '0'; // Declarative program: residual slot + difference/blend operators driven // by the lowering; policy keeps only its scalar decision state. return detail::make_output_diff_program("UCache", seg); @@ -127,9 +125,8 @@ class UCachePolicy final : public ICachePolicy { return {}; } - // ---- Substep interface (ED_CACHE_SUBSTEP). Output method: one substep, reuse - // or compute; decision reuses decide(). ---- - bool supports_substep() const override { return substep_env_on_; } + // ---- Substep interface. Output method: one substep, reuse or compute; + // decision reuses decide(). ---- void set_substep_input(const sd::Tensor* input) override { substep_input_ = input; } void begin_substeps(const StepContext& step, const void* condition_key) override { begin_step(step); @@ -301,7 +298,6 @@ class UCachePolicy final : public ICachePolicy { int consecutive_skipped_steps_ = 0; float accumulated_error_ = 0.0f; int total_active_steps_ = 0; - bool substep_env_on_ = false; bool substep_done_ = false; const void* substep_key_ = nullptr; const sd::Tensor* substep_input_ = nullptr; diff --git a/src/core/optimization/cache/runtime/cache_engine.cpp b/src/core/optimization/cache/runtime/cache_engine.cpp index 90d9eebf..8394f2d0 100644 --- a/src/core/optimization/cache/runtime/cache_engine.cpp +++ b/src/core/optimization/cache/runtime/cache_engine.cpp @@ -74,6 +74,28 @@ bool CacheEngine::init(const ed_sample_params_t& sample_params, return false; } + // Reject TaylorSeer / SenCache on device models (a device store was wired: + // Qwen/Flux). These Feature methods emit FeatureReuse/FeatureCompute, which the + // lowering serves through the host capture/inject hooks — but a device model's + // pipeline routes all Feature-granularity methods into the device-slot branch + // (feature_gpu gate) that sets only the device-slot hooks, not hooks.capture, + // and device runners expose no compute_substep_capture_host. The block-stack + // residual ring therefore never seeds and the method silently degrades to zero + // reuse. MagCache (device-slot single residual) and DiCache (probe) are served by + // their own paths and are unaffected. Disable explicitly rather than run a no-op + // cache that still pays the per-step decision/probe cost. Host models (SD3/Wan, + // no device store) wire the host hooks unconditionally and keep these methods. + if (device_store != nullptr && + (config_.mode == CacheMode::TaylorSeer || config_.mode == CacheMode::SenCache)) { + LOG_WARN("cache disabled: %s is not supported on this model (on-device runner). " + "Its feature-history reuse needs a host capture path that the device " + "pipeline does not provide; use MagCache or DiCache for on-device caching.", + cache_mode_name(config_.mode)); + policy_.reset(); + contract_.reset(); + return false; + } + InferenceConfig inf; inf.config = &config_; inf.sigmas = &sigmas; @@ -82,10 +104,6 @@ bool CacheEngine::init(const ed_sample_params_t& sample_params, contract_->schema().family == ModelFamily::WanVideo; program_ = policy_->compile(contract_->schema(), contract_->topology(), inf); - // Tell the policy whether an on-device metric path exists this run (a device - // store was wired). DiCache gates its substep path on this so a host-only model - // (Wan, no store) keeps its verified legacy host path. - policy_->set_substep_device_available(device_store != nullptr); // Wire the device store BEFORE initialize() so device_backed slots allocate // on-device. Null store (CPU/SP/mmdit/wan) leaves every slot host-backed. state_.set_device_store(device_store); diff --git a/src/core/optimization/cache/state/cache_device_store.hpp b/src/core/optimization/cache/state/cache_device_store.hpp index 30772a31..7000feab 100644 --- a/src/core/optimization/cache/state/cache_device_store.hpp +++ b/src/core/optimization/cache/state/cache_device_store.hpp @@ -22,7 +22,7 @@ class ICacheDeviceStore { // Allocate (or reuse, if shape matches) a persistent device tensor for one // ring entry, keyed by the manager's folded (condition_key<<16 ^ slot) key // plus the ring index. Persists across steps in a context OUTSIDE the per-step - // compute buffer, so it survives run_cache_pass's reset_graph_cut_run_cache(). + // compute buffer, so it survives the substep pass's reset_graph_cut_run_cache(). // Returns an opaque device-tensor handle owned by the store, or nullptr on // failure. `shape` is the block-stack residual extents [ne0, ne1, ne2...]. virtual void* ensure_entry(uint64_t ring_key, int ring_index, diff --git a/src/dit_models/diffusion_model.hpp b/src/dit_models/diffusion_model.hpp index 9728f6a3..e99e5936 100644 --- a/src/dit_models/diffusion_model.hpp +++ b/src/dit_models/diffusion_model.hpp @@ -69,31 +69,8 @@ struct DiffusionModel { // True when the feature seam can run this configuration (plain compute path // only). Models that support feature cache forward this to their runner. virtual bool feature_cache_available() const { return false; } - // Full compute that also captures the cached region's residual. The region - // is the block interval [region_start, region_end); the default (0, -1) is - // the whole block stack. - virtual DiffusionCacheResult compute_capture(int /*n_threads*/, - const DiffusionParams& /*params*/, - int /*region_start*/ = 0, - int /*region_end*/ = -1) { - return {}; - } - // Skip the region's blocks, injecting `feature` as that region's residual. - virtual sd::Tensor compute_inject(int /*n_threads*/, - const DiffusionParams& /*params*/, - const sd::Tensor& /*feature*/, - int /*region_start*/ = 0, - int /*region_end*/ = -1) { - return {}; - } - // Run only the first `probe_depth` blocks; return (before, probe) states. - virtual DiffusionCacheResult compute_probe(int /*n_threads*/, - const DiffusionParams& /*params*/, - int /*probe_depth*/) { - return {}; - } - // Substep-path (ED_CACHE_SUBSTEP) tap-driven host variants — models without a + // Substep-path tap-driven host variants — models without a // device slot (Wan). capture: full forward, host residual (ModelOut-ModelIn) // via taps; probe: shallow prefix, host before/probe via taps. Default no-op. virtual DiffusionCacheResult compute_substep_capture_host(int /*n_threads*/, @@ -106,7 +83,7 @@ struct DiffusionModel { return {}; } // Tap-driven host inject (reuse): x_before + feature over [start,end), region's - // blocks skipped. No CacheGraphScope. Default no-op. + // blocks skipped. Default no-op. virtual sd::Tensor compute_substep_inject_host(int /*n_threads*/, const DiffusionParams& /*params*/, const sd::Tensor& /*feature*/, @@ -264,26 +241,6 @@ struct MMDiTModel : public DiffusionModel { bool supports_feature_cache() const override { return true; } bool feature_cache_available() const override { return mmdit.feature_cache_available(); } - DiffusionCacheResult compute_capture(int n_threads, const DiffusionParams& p, - int region_start = 0, int region_end = -1) override { - return mmdit.compute_capture(n_threads, *p.x, *p.timesteps, - tensor_or_empty(p.context), tensor_or_empty(p.y), - region_start, region_end); - } - - sd::Tensor compute_inject(int n_threads, const DiffusionParams& p, - const sd::Tensor& feature, - int region_start = 0, int region_end = -1) override { - return mmdit.compute_inject(n_threads, *p.x, *p.timesteps, - tensor_or_empty(p.context), tensor_or_empty(p.y), feature, - region_start, region_end); - } - - DiffusionCacheResult compute_probe(int n_threads, const DiffusionParams& p, int probe_depth) override { - return mmdit.compute_probe(n_threads, *p.x, *p.timesteps, - tensor_or_empty(p.context), tensor_or_empty(p.y), probe_depth); - } - DiffusionCacheResult compute_substep_capture_host(int n_threads, const DiffusionParams& p) override { return mmdit.compute_substep_capture(n_threads, *p.x, *p.timesteps, tensor_or_empty(p.context), tensor_or_empty(p.y)); @@ -380,36 +337,6 @@ struct FluxModel : public DiffusionModel { bool supports_feature_cache() const override { return true; } bool feature_cache_available() const override { return flux.feature_cache_available(); } - - DiffusionCacheResult compute_capture(int n_threads, const DiffusionParams& p, - int region_start = 0, int region_end = -1) override { - static const std::vector> empty_ref_latents; - return flux.compute_capture(n_threads, *p.x, *p.timesteps, - tensor_or_empty(p.context), tensor_or_empty(p.c_concat), - tensor_or_empty(p.y), tensor_or_empty(p.guidance), - p.ref_latents ? *p.ref_latents : empty_ref_latents, - p.increase_ref_index, region_start, region_end); - } - - sd::Tensor compute_inject(int n_threads, const DiffusionParams& p, - const sd::Tensor& feature, - int region_start = 0, int region_end = -1) override { - static const std::vector> empty_ref_latents; - return flux.compute_inject(n_threads, *p.x, *p.timesteps, - tensor_or_empty(p.context), tensor_or_empty(p.c_concat), - tensor_or_empty(p.y), tensor_or_empty(p.guidance), - p.ref_latents ? *p.ref_latents : empty_ref_latents, - p.increase_ref_index, feature, region_start, region_end); - } - - DiffusionCacheResult compute_probe(int n_threads, const DiffusionParams& p, int probe_depth) override { - static const std::vector> empty_ref_latents; - return flux.compute_probe(n_threads, *p.x, *p.timesteps, - tensor_or_empty(p.context), tensor_or_empty(p.c_concat), - tensor_or_empty(p.y), tensor_or_empty(p.guidance), - p.ref_latents ? *p.ref_latents : empty_ref_latents, - p.increase_ref_index, probe_depth); - } }; struct AnimaModel : public DiffusionModel { @@ -570,26 +497,6 @@ struct WanModel : public DiffusionModel { bool supports_feature_cache() const override { return true; } bool feature_cache_available() const override { return wan.feature_cache_available(); } - DiffusionCacheResult compute_capture(int n_threads, const DiffusionParams& p, - int region_start = 0, int region_end = -1) override { - return wan.compute_capture(n_threads, *p.x, *p.timesteps, tensor_or_empty(p.context), - tensor_or_empty(p.y), tensor_or_empty(p.c_concat), - region_start, region_end); - } - - sd::Tensor compute_inject(int n_threads, const DiffusionParams& p, - const sd::Tensor& feature, - int region_start = 0, int region_end = -1) override { - return wan.compute_inject(n_threads, *p.x, *p.timesteps, tensor_or_empty(p.context), - tensor_or_empty(p.y), tensor_or_empty(p.c_concat), feature, - region_start, region_end); - } - - DiffusionCacheResult compute_probe(int n_threads, const DiffusionParams& p, int probe_depth) override { - return wan.compute_probe(n_threads, *p.x, *p.timesteps, tensor_or_empty(p.context), - tensor_or_empty(p.y), tensor_or_empty(p.c_concat), probe_depth); - } - DiffusionCacheResult compute_substep_capture_host(int n_threads, const DiffusionParams& p) override { return wan.compute_substep_capture(n_threads, *p.x, *p.timesteps, tensor_or_empty(p.context), tensor_or_empty(p.y), tensor_or_empty(p.c_concat)); @@ -683,29 +590,6 @@ struct QwenImageModel : public DiffusionModel { bool supports_feature_cache() const override { return true; } bool feature_cache_available() const override { return qwen_image.feature_cache_available(); } - - DiffusionCacheResult compute_capture(int n_threads, const DiffusionParams& p, - int region_start = 0, int region_end = -1) override { - static const std::vector> empty_ref_latents; - return qwen_image.compute_capture(n_threads, *p.x, *p.timesteps, tensor_or_empty(p.context), - p.ref_latents ? *p.ref_latents : empty_ref_latents, true, - region_start, region_end); - } - - sd::Tensor compute_inject(int n_threads, const DiffusionParams& p, - const sd::Tensor& feature, - int region_start = 0, int region_end = -1) override { - static const std::vector> empty_ref_latents; - return qwen_image.compute_inject(n_threads, *p.x, *p.timesteps, tensor_or_empty(p.context), - p.ref_latents ? *p.ref_latents : empty_ref_latents, true, feature, - region_start, region_end); - } - - DiffusionCacheResult compute_probe(int n_threads, const DiffusionParams& p, int probe_depth) override { - static const std::vector> empty_ref_latents; - return qwen_image.compute_probe(n_threads, *p.x, *p.timesteps, tensor_or_empty(p.context), - p.ref_latents ? *p.ref_latents : empty_ref_latents, true, probe_depth); - } }; struct ZImageModel : public DiffusionModel { diff --git a/src/dit_models/models/flux.hpp b/src/dit_models/models/flux.hpp index dee7ba73..f7fe2918 100644 --- a/src/dit_models/models/flux.hpp +++ b/src/dit_models/models/flux.hpp @@ -3577,29 +3577,15 @@ namespace Flux { const int64_t flux_full_img_seq = img->ne[1]; const int64_t flux_full_txt_seq = txt->ne[1]; // Cache seam: the block stack transforms the image stream `img`. The - // cached quantity is that stream's residual. Disabled under SP - // (block-loop tensors are sequence-sharded, so before/after differ - // in shape). See cache_graph_scope.hpp. - sd::CacheGraphScope* cache_scope = flux_sp_enabled(ctx) ? nullptr : ctx->cache_scope; - // A sub-region (finite region_end) is cached inside the double-block - // loop, where `img` is coherent. The default whole-stack region spans - // the double AND single loops and is captured after the streams are - // recombined and `img` is sliced back out (the original behaviour). - const bool cache_subregion = cache_scope != nullptr && cache_scope->is_subregion(); - const bool cache_inject = cache_scope != nullptr && cache_scope->inject_mode(); - const bool cache_whole_inject = cache_inject && !cache_subregion; + // cached quantity is that stream's residual. // Tap-driven whole-stack inject (substep reuse): skip both block loops and - // reconstruct x_before + residual via the registry (no CacheGraphScope). + // reconstruct x_before + residual via the registry. const bool tap_whole_inject = ctx->tap_registry != nullptr && ctx->tap_registry->inject_active(); - const bool whole_inject = cache_whole_inject || tap_whole_inject; + const bool whole_inject = tap_whole_inject; ggml_tensor* cache_img_before = img; - if (cache_scope != nullptr && !cache_subregion) { - cache_scope->snapshot_before(cache_img_before); // whole-stack anchor - } // Substep tap: block-stack input anchor (ModelIn). Conditional no-op - // unless the middle layer requested it this substep. Coexists with the - // legacy cache_scope path (double-write during migration). + // unless the middle layer requested it this substep. tap(ctx, edgedit::cache::AnchorRef::model_in(), cache_img_before); bool use_sp_mainline = flux_sp_enabled(ctx); #ifdef ED_DEBUG_SP_COMM @@ -3831,18 +3817,6 @@ namespace Flux { if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i) != skip_layers.end()) { continue; } - // Sub-region seam (inside the double-block loop): inject skips the - // region, capture snapshots its input and builds its residual. - if (cache_subregion && cache_inject) { - if (ggml_tensor* injected = cache_scope->step_inject_region(ctx->ggml_ctx, i, img)) { - img = injected; - i = cache_scope->inject_resume_index(params.depth) - 1; - continue; - } - } - if (cache_subregion) { - cache_scope->begin_region(i, img); - } auto block = std::dynamic_pointer_cast(blocks["double_blocks." + std::to_string(i)]); @@ -3901,22 +3875,6 @@ namespace Flux { if (ctx->tap_registry != nullptr && ctx->tap_registry->stop_after(i)) { return img; } - // Sub-region capture: build the region residual at its last block. - if (cache_subregion) { - cache_scope->end_region(ctx->ggml_ctx, i, params.depth, img); - } - // GPU DiCache: on a full (capture) step, snapshot the probe-depth - // hidden state so the runner refreshes prev_probe for next step. - if (cache_scope != nullptr) { - cache_scope->record_capture_probe_state(ctx->ggml_ctx, i, img); - } - // Probe step: stop after the first probe_depth double blocks and - // hand the shallow image state back to the runner. The final - // layer is not run; the runner only reads the probe/before pair. - if (cache_scope != nullptr && cache_scope->stop_after_block(i)) { - cache_scope->on_probe(img); - return img; - } #ifdef ED_DEBUG_SP_COMM if (debug_compare_double_inner && i >= debug_compare_double_block_idx) { @@ -4161,28 +4119,15 @@ namespace Flux { } #endif - // Cache seam (whole-stack path): on a capture step expose the - // image-stream residual for the runner to read back; on an inject - // step the loops above ran zero blocks so `img` equals - // cache_img_before, and we add the reconstructed residual to - // approximate the full stack output. A sub-region was already - // captured/injected inside the double-block loop above. - if (cache_scope != nullptr && !cache_subregion) { - if (cache_scope->inject_mode()) { - img = cache_scope->add_injected(ctx->ggml_ctx, cache_img_before); - } else if (cache_scope->capture_mode()) { - cache_scope->build_feature(ctx->ggml_ctx, img); - } - } - // Tap-driven whole-stack inject (substep reuse): the loops ran zero - // blocks, so reconstruct x_before + residual via the registry. No - // CacheGraphScope. + // Cache seam (whole-stack path): tap-driven whole-stack inject + // (substep reuse): the loops ran zero blocks, so reconstruct + // x_before + residual via the registry. if (tap_whole_inject) { img = build_tap_inject(ctx, cache_img_before); } // Substep tap: block-stack output anchor (ModelOut) — the residual's - // "after" point (post-recombine, before final_layer), matching - // build_feature's capture site. Conditional no-op unless requested. + // "after" point (post-recombine, before final_layer). Conditional no-op + // unless requested. tap(ctx, edgedit::cache::AnchorRef::model_out(), img); if (final_layer) { @@ -4899,134 +4844,8 @@ namespace Flux { dicache_gpu_states_.clear(); } - // Ensure per-branch persistent state sized to the just-computed named - // cache tensors (feature/probe/before all share the block-stack shape). - DiCacheGpuState& ensure_dicache_gpu_state_from_named(const void* branch_key) { - ggml_tensor* feat = get_cache_tensor_by_name(kCacheFeatureName); - if (feat == nullptr) { - return dicache_gpu_states_[branch_key]; // empty; caller checks buffer - } - return ensure_dicache_gpu_state(branch_key, feat->ne[0], feat->ne[1], feat->ne[2]); - } - int dicache_probe_depth_ = 1; // set by the pipeline before capture/probe - sd::DiffusionCacheResult compute_capture(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context, - const sd::Tensor& c_concat, - const sd::Tensor& y, - const sd::Tensor& guidance, - const std::vector>& ref_latents, - bool increase_ref_index, - int region_start = 0, - int region_end = -1, - const void* branch_key = nullptr) { - sd::CacheGraphScope scope; - scope.mode = sd::CacheGraphScope::Mode::Capture; - scope.region_start = region_start; - scope.region_end = region_end; - const bool gpu = dicache_gpu_enabled() && branch_key != nullptr; - if (gpu) { - scope.gpu_metric = true; // record probe-depth state during capture - scope.probe_depth = std::max(1, dicache_probe_depth_); - } - std::function handoff = nullptr; - if (gpu) { - handoff = [&]() { - // Refresh the persistent cross-step state device-to-device from - // the just-computed named tensors, with a 2-deep ping-pong ring. - DiCacheGpuState& s = ensure_dicache_gpu_state_from_named(branch_key); - if (s.buffer == nullptr) return; - // full residual ring: prev2 <- prev1, prev1 <- feature - std::swap(s.resid_prev1, s.resid_prev2); - copy_named_cache_tensor_to(kCacheFeatureName, s.resid_prev1); - s.resid_count = std::min(2, s.resid_count + 1); - // probe residual ring: prev2 <- prev1, prev1 <- (probe - before) - std::swap(s.probe_prev1, s.probe_prev2); - copy_named_cache_tensor_to(kCacheProbeResidName, s.probe_prev1); - s.probe_resid_count = std::min(2, s.probe_resid_count + 1); - // raw probe & input for the delta_y / delta_x decision metrics - copy_named_cache_tensor_to(kCacheProbeName, s.prev_probe); - copy_named_cache_tensor_to(kCacheBeforeName, s.prev_input); - s.has_probe_hist = true; - }; - } - auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, {}); - }; - auto pass = run_cache_pass(get_graph, n_threads, &scope, x.dim(), handoff); - sd::DiffusionCacheResult out; - out.output = std::move(pass.output); - out.feature = std::move(pass.feature); - return out; - } - - sd::Tensor compute_inject(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context, - const sd::Tensor& c_concat, - const sd::Tensor& y, - const sd::Tensor& guidance, - const std::vector>& ref_latents, - bool increase_ref_index, - const sd::Tensor& feature, - int region_start = 0, - int region_end = -1) { - sd::CacheGraphScope scope; - scope.mode = sd::CacheGraphScope::Mode::Inject; - scope.region_start = region_start; - scope.region_end = region_end; - inject_feature_host_ = feature; - auto get_graph = [&]() -> ggml_cgraph* { - scope.inject_feature = make_input(inject_feature_host_); - return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, {}); - }; - auto pass = run_cache_pass(get_graph, n_threads, &scope, x.dim()); - return std::move(pass.output); - } - - // GPU DiCache reuse step: reconstruct the residual on-device from the - // persistent full-residual ring and a host-clamped gamma, injecting - // x_before + resid_prev2 + gamma*(resid_prev1 - resid_prev2) - // No ~50MB host residual is built or uploaded. Returns empty if state is - // not ready (caller falls back to a capturing full step). - sd::Tensor compute_inject_gpu(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context, - const sd::Tensor& c_concat, - const sd::Tensor& y, - const sd::Tensor& guidance, - const std::vector>& ref_latents, - bool increase_ref_index, - float gamma, - const void* branch_key, - int region_start = 0, - int region_end = -1) { - auto it = dicache_gpu_states_.find(branch_key); - if (it == dicache_gpu_states_.end() || it->second.buffer == nullptr || - it->second.resid_count < 2) { - return {}; // not enough history yet - } - DiCacheGpuState& s = it->second; - sd::CacheGraphScope scope; - scope.mode = sd::CacheGraphScope::Mode::Inject; - scope.region_start = region_start; - scope.region_end = region_end; - scope.resid_prev1_node = s.resid_prev1; - scope.resid_prev2_node = s.resid_prev2; - gamma_scalar_host_ = sd::Tensor({1}, std::vector{gamma}); - auto get_graph = [&]() -> ggml_cgraph* { - scope.gamma_scalar = make_input(gamma_scalar_host_); - return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, {}); - }; - auto pass = run_cache_pass(get_graph, n_threads, &scope, x.dim()); - return std::move(pass.output); - } - // ---- Feature-granularity on-GPU reuse (MagCache / TaylorSeer-single) ---- // MagCache's host reuse pays a ~50MB reconstruct copy + a 2nd host copy + // an H2D upload per skip step. When enabled, the last captured block-stack @@ -5043,81 +4862,11 @@ namespace Flux { // ---- Declarative device-slot seam (B2): on-device single-residual reuse // where the persistent residual tensor is a - // CacheStateManager device slot (passed in) instead of DiCacheGpuState. - // capture copies the block-stack residual into `slot`; inject reads it. ---- - sd::Tensor compute_capture_to_slot(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context, - const sd::Tensor& c_concat, - const sd::Tensor& y, - const sd::Tensor& guidance, - const std::vector>& ref_latents, - bool increase_ref_index, - const std::function&)>& alloc_slot, - int region_start, - int region_end) { - sd::CacheGraphScope scope; - scope.mode = sd::CacheGraphScope::Mode::Capture; - scope.region_start = region_start; - scope.region_end = region_end; - std::function handoff = [&]() { - // The residual named tensor is resident; allocate the device slot at - // its true (packed block-stack seq) shape, then d2d-copy into it. - ggml_tensor* feat = get_cache_tensor_by_name(kCacheFeatureName); - if (feat == nullptr) { - return; - } - // Copy feat's exact dims so the slot is byte- and broadcast-identical - // (add_injected does ggml_add(x_before, slot)). ggml_n_dims collapses - // trailing singletons but never a size-1 inner dim. - std::vector shape; - const int nd = std::max(1, ggml_n_dims(feat)); - for (int i = 0; i < nd; ++i) { - shape.push_back(feat->ne[i]); - } - ggml_tensor* slot = static_cast(alloc_slot(shape)); - if (slot != nullptr && ggml_nbytes(slot) == ggml_nbytes(feat)) { - copy_named_cache_tensor_to(kCacheFeatureName, slot); - } - }; - auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, {}); - }; - auto pass = run_cache_pass(get_graph, n_threads, &scope, x.dim(), handoff); - return std::move(pass.output); - } + // CacheStateManager device slot (passed in) instead of DiCacheGpuState. ---- - sd::Tensor compute_inject_from_slot(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context, - const sd::Tensor& c_concat, - const sd::Tensor& y, - const sd::Tensor& guidance, - const std::vector>& ref_latents, - bool increase_ref_index, - ggml_tensor* slot, - int region_start, - int region_end) { - if (slot == nullptr) { - return {}; - } - sd::CacheGraphScope scope; - scope.mode = sd::CacheGraphScope::Mode::Inject; - scope.region_start = region_start; - scope.region_end = region_end; - scope.resid_prev1_node = slot; // single residual: gpu_reuse_available() - auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, {}); - }; - auto pass = run_cache_pass(get_graph, n_threads, &scope, x.dim()); - return std::move(pass.output); - } - - // ---- Substep-path (ED_CACHE_SUBSTEP) tap-driven device inject. The whole + // ---- Substep-path tap-driven device inject. The whole // stack is skipped (registry inject_active) and x_before + residual is - // reconstructed via build_tap_inject in the forward. No CacheGraphScope. ---- + // reconstructed via build_tap_inject in the forward. ---- sd::Tensor compute_substep_inject_slot(int n_threads, const sd::Tensor& x, const sd::Tensor& timesteps, @@ -5174,60 +4923,9 @@ namespace Flux { return std::move(pass.output); } - sd::DiffusionCacheResult compute_probe(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context, - const sd::Tensor& c_concat, - const sd::Tensor& y, - const sd::Tensor& guidance, - const std::vector>& ref_latents, - bool increase_ref_index, - int probe_depth, - const void* branch_key = nullptr) { - sd::CacheGraphScope scope; - scope.mode = sd::CacheGraphScope::Mode::Probe; - scope.probe_depth = probe_depth; - // GPU DiCache path (env-gated until verified): attach the persistent - // prev_* tensors so build_probe_metrics emits delta_y/delta_x/gamma - // reductions in-graph; only scalars are read back. - DiCacheGpuState* gpu = nullptr; - if (dicache_gpu_enabled() && branch_key != nullptr) { - auto it = dicache_gpu_states_.find(branch_key); - if (it != dicache_gpu_states_.end() && it->second.has_probe_hist) { - gpu = &it->second; - scope.gpu_metric = true; - scope.prev_probe_node = gpu->prev_probe; - scope.prev_input_node = gpu->prev_input; - if (gpu->probe_resid_count >= 2) { - scope.probe_prev1_node = gpu->probe_prev1; - scope.probe_prev2_node = gpu->probe_prev2; - } - } - } - auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, {}); - }; - auto pass = run_cache_pass(get_graph, n_threads, &scope, x.dim()); - sd::DiffusionCacheResult out; - if (scope.gpu_metric) { - // GPU path: no ~50MB host tensors; only scalars. - out.delta_y = pass.delta_y; - out.delta_x = pass.delta_x; - out.gamma = pass.gamma; - return out; - } - // Prefer the explicit probe_node readback; fall back to the graph - // output (which is the probe state) if the node wasn't recorded. - out.probe = pass.probe.empty() ? std::move(pass.output) : std::move(pass.probe); - out.before = std::move(pass.before); - return out; - } - - // ---- Substep-path (ED_CACHE_SUBSTEP) tap-driven capture. Replaces the + // ---- Substep-path tap-driven capture. Replaces the // capture_to_slot hook: requests ModelIn/ModelOut taps, weaves the - // (ModelOut - ModelIn) residual, d2d-copies it into the device slot. No - // CacheGraphScope on the capture path. ---- + // (ModelOut - ModelIn) residual, d2d-copies it into the device slot. ---- sd::Tensor compute_substep_capture(int n_threads, const sd::Tensor& x, const sd::Tensor& timesteps, @@ -5264,9 +4962,9 @@ namespace Flux { return std::move(pass.output); } - // ---- Substep-path (ED_CACHE_SUBSTEP) tap-driven DiCache probe. Requests + // ---- Substep-path tap-driven DiCache probe. Requests // ModelIn + BlockOut[m-1] taps, stops after m double blocks, threads the - // persistent operands, weaves delta_y/delta_x/gamma. No CacheGraphScope. ---- + // persistent operands, weaves delta_y/delta_x/gamma. ---- sd::DiffusionCacheResult compute_substep_probe(int n_threads, const sd::Tensor& x, const sd::Tensor& timesteps, @@ -5320,6 +5018,70 @@ namespace Flux { return out; } + // ---- Substep-path tap-driven DiCache seed capture (device). Mirrors Qwen's + // compute_substep_capture_probe: a full forward whose post-readback refreshes + // the persistent DiCacheGpuState residual/probe rings device-to-device from the + // tap-woven residual/probe nodes. No device slot (the ring is runner-owned, + // not a StateManager slot). ---- + void writeback_from_taps(const void* branch_key) { + ggml_tensor* feat = get_cache_tensor_by_name("ed_cache_feature"); + if (feat == nullptr) { + return; + } + DiCacheGpuState& s = + ensure_dicache_gpu_state(branch_key, feat->ne[0], feat->ne[1], feat->ne[2]); + if (s.buffer == nullptr) { + return; + } + const int m = std::max(1, dicache_probe_depth_); + const std::string probe_tap = "ed_tap:block_out[" + std::to_string(m - 1) + "]"; + // A silent copy failure (shape/name mismatch) would leave the ring stale + // and silently degrade the decision, so surface it loudly. + auto cp = [&](const char* name, ggml_tensor* dst) { + if (!copy_named_cache_tensor_to(name, dst)) { + LOG_ERROR("dicache writeback copy failed: %s", name); + } + }; + // full residual ring: prev2 <- prev1, prev1 <- feature + std::swap(s.resid_prev1, s.resid_prev2); + cp("ed_cache_feature", s.resid_prev1); + s.resid_count = std::min(2, s.resid_count + 1); + // probe residual ring: prev2 <- prev1, prev1 <- (probe - before) + std::swap(s.probe_prev1, s.probe_prev2); + cp("ed_cache_probe_resid", s.probe_prev1); + s.probe_resid_count = std::min(2, s.probe_resid_count + 1); + // raw probe & input for the delta_y / delta_x decision metrics + cp(probe_tap.c_str(), s.prev_probe); + cp("ed_tap:model_in", s.prev_input); + s.has_probe_hist = true; + } + + sd::Tensor compute_substep_capture_probe(int n_threads, + const sd::Tensor& x, + const sd::Tensor& timesteps, + const sd::Tensor& context, + const sd::Tensor& c_concat, + const sd::Tensor& y, + const sd::Tensor& guidance, + const std::vector>& ref_latents, + bool increase_ref_index, + int probe_depth, + const void* branch_key) { + const int m = std::max(1, probe_depth); + edgedit::cache::TapRegistry reg; + reg.set_requested({edgedit::cache::AnchorRef::model_in(), + edgedit::cache::AnchorRef::block_out(m - 1), + edgedit::cache::AnchorRef::model_out()}); + reg.set_capture_residual(true); // weave ed_cache_feature = ModelOut - ModelIn + reg.set_capture_writeback(m); // weave ed_cache_probe_resid = block_out(m-1) - ModelIn + std::function handoff = [&]() { writeback_from_taps(branch_key); }; + auto get_graph = [&]() -> ggml_cgraph* { + return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, {}); + }; + auto pass = run_substep_pass(get_graph, n_threads, ®, x.dim(), {}, handoff); + return std::move(pass.output); + } + // GPU DiCache is the default: on-device metric + reconstruction avoids the // ~4.5s/img host scalar work and the ~50MB/step GPU->host readback, and is // slightly MORE accurate (exact GPU reductions). Set ED_DICACHE_GPU=0 to diff --git a/src/dit_models/models/mmdit.hpp b/src/dit_models/models/mmdit.hpp index 3a9d8756..a8c83056 100644 --- a/src/dit_models/models/mmdit.hpp +++ b/src/dit_models/models/mmdit.hpp @@ -3274,14 +3274,10 @@ struct MMDiT : public GGMLBlock { auto final_layer = std::dynamic_pointer_cast(blocks["final_layer"]); // Cache seam: the joint block stack transforms the image stream `x`. - // Disabled under SP (block-loop tensors are sequence-sharded). The cached - // region is blocks [region_start, region_end); the default whole-stack - // region matches the pre-region behaviour. See cache_graph_scope.hpp. - sd::CacheGraphScope* cache_scope = use_sp_mainline ? nullptr : ctx->cache_scope; - const bool cache_inject = cache_scope != nullptr && cache_scope->inject_mode(); - + // The cached region is blocks [region_start, region_end); the default + // whole-stack region matches the pre-region behaviour. // Substep tap: block-stack input anchor (ModelIn). Conditional no-op unless - // requested. Coexists with the legacy cache_scope path (double-write). + // requested. tap(ctx, edgedit::cache::AnchorRef::model_in(), x); ggml_tensor* cache_x_before = x; // x_before for tap-driven inject @@ -3293,24 +3289,12 @@ struct MMDiT : public GGMLBlock { // Tap-driven inject (substep reuse): at the region start, replace the // stream with the reconstructed x_before + residual and jump past the - // region — no CacheGraphScope. + // region. if (ctx->tap_registry != nullptr && ctx->tap_registry->inject_at(i)) { x = build_tap_inject(ctx, cache_x_before); i = ctx->tap_registry->inject_resume() - 1; continue; } - // Inject: at the region start, collapse the region to a residual add - // and jump past its blocks (blocks before region_start still run). - if (cache_inject) { - if (ggml_tensor* injected = cache_scope->step_inject_region(ctx->ggml_ctx, i, x)) { - x = injected; - i = cache_scope->inject_resume_index(depth) - 1; // ++ lands at resume - continue; - } - } - if (cache_scope != nullptr) { - cache_scope->begin_region(i, x); - } auto block = std::dynamic_pointer_cast(blocks["joint_blocks." + std::to_string(i)]); @@ -3337,13 +3321,6 @@ struct MMDiT : public GGMLBlock { if (ctx->tap_registry != nullptr && ctx->tap_registry->stop_after(i)) { return x; } - if (cache_scope != nullptr) { - cache_scope->end_region(ctx->ggml_ctx, i, depth, x); - if (cache_scope->stop_after_block(i)) { - cache_scope->on_probe(x); - return x; - } - } } // Substep tap: block-stack output anchor (ModelOut) — the residual's "after" @@ -3536,68 +3513,10 @@ struct MMDiTRunner : public GGMLRunner { return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } - // ---- Cache seam passes (Feature/Probe policies). ---- - sd::DiffusionCacheResult compute_capture(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context, - const sd::Tensor& y, - int region_start = 0, - int region_end = -1) { - sd::CacheGraphScope scope; - scope.mode = sd::CacheGraphScope::Mode::Capture; - scope.region_start = region_start; - scope.region_end = region_end; - auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, y, {}); }; - auto pass = run_cache_pass(get_graph, n_threads, &scope, x.dim()); - sd::DiffusionCacheResult out; - out.output = std::move(pass.output); - out.feature = std::move(pass.feature); - return out; - } - - sd::Tensor compute_inject(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context, - const sd::Tensor& y, - const sd::Tensor& feature, - int region_start = 0, - int region_end = -1) { - sd::CacheGraphScope scope; - scope.mode = sd::CacheGraphScope::Mode::Inject; - scope.region_start = region_start; - scope.region_end = region_end; - inject_feature_host_ = feature; - auto get_graph = [&]() -> ggml_cgraph* { - scope.inject_feature = make_input(inject_feature_host_); - return build_graph(x, timesteps, context, y, {}); - }; - auto pass = run_cache_pass(get_graph, n_threads, &scope, x.dim()); - return std::move(pass.output); - } - - sd::DiffusionCacheResult compute_probe(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context, - const sd::Tensor& y, - int probe_depth) { - sd::CacheGraphScope scope; - scope.mode = sd::CacheGraphScope::Mode::Probe; - scope.probe_depth = probe_depth; - auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, y, {}); }; - auto pass = run_cache_pass(get_graph, n_threads, &scope, x.dim()); - sd::DiffusionCacheResult out; - out.probe = pass.probe.empty() ? std::move(pass.output) : std::move(pass.probe); - out.before = std::move(pass.before); - return out; - } - - // ---- Substep-path (ED_CACHE_SUBSTEP) tap-driven host cache passes. SD3/mmdit + // ---- Substep-path tap-driven host cache passes. SD3/mmdit // has no device slot, so residual/probe tensors are read back to host. Capture // weaves (ModelOut - ModelIn); probe reads before/probe; inject reconstructs - // x_before + host feature via the forward's registry inject. No CacheGraphScope. ---- + // x_before + host feature via the forward's registry inject. ---- sd::DiffusionCacheResult compute_substep_capture(int n_threads, const sd::Tensor& x, const sd::Tensor& timesteps, diff --git a/src/dit_models/models/qwen_image.hpp b/src/dit_models/models/qwen_image.hpp index fc8e5b0e..41b4dce2 100644 --- a/src/dit_models/models/qwen_image.hpp +++ b/src/dit_models/models/qwen_image.hpp @@ -2966,40 +2966,22 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* // sd::ggml_graph_cut::mark_graph_cut(t_emb, "qwen_image.prelude", "t_emb"); // Cache seam: the transformer stack transforms the image stream - // `img`. Disabled under SP (block-loop tensors are sequence-sharded). - // The cached region is blocks [region_start, region_end); the default - // whole-stack region matches the pre-region behaviour. Capture builds - // the region residual in-loop at the region's last block, so the seam - // is only active on the non-SP path (SP is gated off above) where no - // post-loop gather rewrites `img`. - sd::CacheGraphScope* cache_scope = use_sp_mainline ? nullptr : ctx->cache_scope; - const bool cache_inject = cache_scope != nullptr && cache_scope->inject_mode(); - + // `img`. The cached region is blocks [region_start, region_end); the + // default whole-stack region matches the pre-region behaviour. // Substep tap: block-stack input anchor (ModelIn). Conditional — a no-op - // unless the middle layer requested it this substep. Coexists with the - // legacy cache_scope path above (double-write during migration). + // unless the middle layer requested it this substep. tap(ctx, edgedit::cache::AnchorRef::model_in(), img); ggml_tensor* cache_img_before = img; // x_before for tap-driven inject for (int i = 0; i < params.num_layers; i++) { // Tap-driven inject (substep reuse): at the region start, replace the // stream with the reconstructed x_before + residual and jump past the - // region — no CacheGraphScope. + // region. if (ctx->tap_registry != nullptr && ctx->tap_registry->inject_at(i)) { img = build_tap_inject(ctx, cache_img_before); i = ctx->tap_registry->inject_resume() - 1; continue; } - if (cache_inject) { - if (ggml_tensor* injected = cache_scope->step_inject_region(ctx->ggml_ctx, i, img)) { - img = injected; - i = cache_scope->inject_resume_index(params.num_layers) - 1; - continue; - } - } - if (cache_scope != nullptr) { - cache_scope->begin_region(i, img); - } auto block = std::dynamic_pointer_cast(blocks["transformer_blocks." + std::to_string(i)]); @@ -3032,17 +3014,6 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* if (ctx->tap_registry != nullptr && ctx->tap_registry->stop_after(i)) { return img; } - if (cache_scope != nullptr) { - cache_scope->end_region(ctx->ggml_ctx, i, params.num_layers, img); - // GPU DiCache: on a full (capture) step, snapshot the - // probe-depth hidden state so the runner refreshes prev_probe - // for the next step's decision. - cache_scope->record_capture_probe_state(ctx->ggml_ctx, i, img); - if (cache_scope->stop_after_block(i)) { - cache_scope->on_probe(img); - return img; - } - } } // Substep tap: block-stack output anchor (ModelOut). This is the @@ -3201,16 +3172,6 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* dicache_gpu_states_.clear(); } - // Ensure per-branch persistent state sized to the just-computed named cache - // tensors (feature/probe/before all share the block-stack shape). - DiCacheGpuState& ensure_dicache_gpu_state_from_named(const void* branch_key) { - ggml_tensor* feat = get_cache_tensor_by_name(kCacheFeatureName); - if (feat == nullptr) { - return dicache_gpu_states_[branch_key]; // empty; caller checks buffer - } - return ensure_dicache_gpu_state(branch_key, feat->ne[0], feat->ne[1], feat->ne[2]); - } - QwenImageRunner(ggml_backend_t backend, bool offload_params_to_cpu, const String2TensorStorage& tensor_storage_map = {}, @@ -3368,114 +3329,6 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } - // ---- Cache seam passes (Feature/Probe policies). ---- - sd::DiffusionCacheResult compute_capture(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context, - const std::vector>& ref_latents, - bool increase_ref_index, - int region_start = 0, - int region_end = -1, - const void* branch_key = nullptr) { - sd::CacheGraphScope scope; - scope.mode = sd::CacheGraphScope::Mode::Capture; - scope.region_start = region_start; - scope.region_end = region_end; - const bool gpu = dicache_gpu_enabled() && branch_key != nullptr; - if (gpu) { - scope.gpu_metric = true; // record probe-depth state during capture - scope.probe_depth = std::max(1, dicache_probe_depth_); - } - std::function handoff = nullptr; - if (gpu) { - handoff = [&]() { - // Refresh the persistent cross-step state device-to-device from - // the just-computed named tensors, with a 2-deep ping-pong ring. - DiCacheGpuState& s = ensure_dicache_gpu_state_from_named(branch_key); - if (s.buffer == nullptr) return; - // full residual ring: prev2 <- prev1, prev1 <- feature - std::swap(s.resid_prev1, s.resid_prev2); - copy_named_cache_tensor_to(kCacheFeatureName, s.resid_prev1); - s.resid_count = std::min(2, s.resid_count + 1); - // probe residual ring: prev2 <- prev1, prev1 <- (probe - before) - std::swap(s.probe_prev1, s.probe_prev2); - copy_named_cache_tensor_to(kCacheProbeResidName, s.probe_prev1); - s.probe_resid_count = std::min(2, s.probe_resid_count + 1); - // raw probe & input for the delta_y / delta_x decision metrics - copy_named_cache_tensor_to(kCacheProbeName, s.prev_probe); - copy_named_cache_tensor_to(kCacheBeforeName, s.prev_input); - s.has_probe_hist = true; - }; - } - auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(x, timesteps, context, ref_latents, increase_ref_index); - }; - auto pass = run_cache_pass(get_graph, n_threads, &scope, x.dim(), handoff); - sd::DiffusionCacheResult out; - out.output = std::move(pass.output); - out.feature = std::move(pass.feature); - return out; - } - - sd::Tensor compute_inject(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context, - const std::vector>& ref_latents, - bool increase_ref_index, - const sd::Tensor& feature, - int region_start = 0, - int region_end = -1) { - sd::CacheGraphScope scope; - scope.mode = sd::CacheGraphScope::Mode::Inject; - scope.region_start = region_start; - scope.region_end = region_end; - inject_feature_host_ = feature; - auto get_graph = [&]() -> ggml_cgraph* { - scope.inject_feature = make_input(inject_feature_host_); - return build_graph(x, timesteps, context, ref_latents, increase_ref_index); - }; - auto pass = run_cache_pass(get_graph, n_threads, &scope, x.dim()); - return std::move(pass.output); - } - - // GPU DiCache reuse step: reconstruct the residual on-device from the - // persistent full-residual ring and a host-clamped gamma, injecting - // x_before + resid_prev2 + gamma*(resid_prev1 - resid_prev2) - // No ~50MB host residual is built or uploaded. Returns empty if state is - // not ready (caller falls back to a capturing full step). - sd::Tensor compute_inject_gpu(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context, - const std::vector>& ref_latents, - bool increase_ref_index, - float gamma, - const void* branch_key, - int region_start = 0, - int region_end = -1) { - auto it = dicache_gpu_states_.find(branch_key); - if (it == dicache_gpu_states_.end() || it->second.buffer == nullptr || - it->second.resid_count < 2) { - return {}; // not enough history yet - } - DiCacheGpuState& s = it->second; - sd::CacheGraphScope scope; - scope.mode = sd::CacheGraphScope::Mode::Inject; - scope.region_start = region_start; - scope.region_end = region_end; - scope.resid_prev1_node = s.resid_prev1; - scope.resid_prev2_node = s.resid_prev2; - gamma_scalar_host_ = sd::Tensor({1}, std::vector{gamma}); - auto get_graph = [&]() -> ggml_cgraph* { - scope.gamma_scalar = make_input(gamma_scalar_host_); - return build_graph(x, timesteps, context, ref_latents, increase_ref_index); - }; - auto pass = run_cache_pass(get_graph, n_threads, &scope, x.dim()); - return std::move(pass.output); - } - // ---- Feature-granularity on-GPU reuse (MagCache / TaylorSeer-single) ---- // When enabled, the last captured block-stack residual stays resident // on-device and is injected as x_before + residual on skips, avoiding the @@ -3490,77 +3343,8 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* return v[0] != '0'; } - // ---- Declarative device-slot seam (B2): same on-device single-residual - // reuse as compute_*_feature_gpu, but the persistent residual tensor is a - // CacheStateManager device slot (passed in) instead of DiCacheGpuState. - // capture copies the block-stack residual into `slot`; inject reads it. ---- - sd::Tensor compute_capture_to_slot(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context, - const std::vector>& ref_latents, - bool increase_ref_index, - const std::function&)>& alloc_slot, - int region_start, - int region_end) { - sd::CacheGraphScope scope; - scope.mode = sd::CacheGraphScope::Mode::Capture; - scope.region_start = region_start; - scope.region_end = region_end; - std::function handoff = [&]() { - // The residual named tensor is resident; allocate the device slot at - // its true (packed block-stack seq) shape, then d2d-copy into it. - ggml_tensor* feat = get_cache_tensor_by_name(kCacheFeatureName); - if (feat == nullptr) { - return; - } - // Copy feat's exact dims so the slot is byte- and broadcast-identical - // (add_injected does ggml_add(x_before, slot)). ggml_n_dims collapses - // trailing singletons but never a size-1 inner dim. - std::vector shape; - const int nd = std::max(1, ggml_n_dims(feat)); - for (int i = 0; i < nd; ++i) { - shape.push_back(feat->ne[i]); - } - ggml_tensor* slot = static_cast(alloc_slot(shape)); - if (slot != nullptr && ggml_nbytes(slot) == ggml_nbytes(feat)) { - copy_named_cache_tensor_to(kCacheFeatureName, slot); - } - }; - auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(x, timesteps, context, ref_latents, increase_ref_index); - }; - auto pass = run_cache_pass(get_graph, n_threads, &scope, x.dim(), handoff); - return std::move(pass.output); - } - - sd::Tensor compute_inject_from_slot(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context, - const std::vector>& ref_latents, - bool increase_ref_index, - ggml_tensor* slot, - int region_start, - int region_end) { - if (slot == nullptr) { - return {}; - } - sd::CacheGraphScope scope; - scope.mode = sd::CacheGraphScope::Mode::Inject; - scope.region_start = region_start; - scope.region_end = region_end; - scope.resid_prev1_node = slot; // single residual: gpu_reuse_available() - auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(x, timesteps, context, ref_latents, increase_ref_index); - }; - auto pass = run_cache_pass(get_graph, n_threads, &scope, x.dim()); - return std::move(pass.output); - } - - // ---- Substep-path (ED_CACHE_SUBSTEP) tap-driven device inject. MagCache: - // x_before + slot; the forward's registry inject reconstructs it. No - // CacheGraphScope. ---- + // ---- Substep-path tap-driven device inject. MagCache: + // x_before + slot; the forward's registry inject reconstructs it. ---- sd::Tensor compute_substep_inject_slot(int n_threads, const sd::Tensor& x, const sd::Tensor& timesteps, @@ -3584,7 +3368,7 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* } // DiCache device reuse: x_before + resid2 + gamma*(resid1-resid2) from the - // persistent 2-deep ring, gamma uploaded as a [1] input. No CacheGraphScope. + // persistent 2-deep ring, gamma uploaded as a [1] input. sd::Tensor compute_substep_inject_gpu(int n_threads, const sd::Tensor& x, const sd::Tensor& timesteps, @@ -3613,11 +3397,10 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* return std::move(pass.output); } - // ---- Substep-path (ED_CACHE_SUBSTEP) tap-driven capture. Replaces the + // ---- Substep-path tap-driven capture. Replaces the // capture_to_slot hook: sets a TapRegistry requesting ModelIn/ModelOut, // asks build_graph to weave the (ModelOut-ModelIn) residual, runs the pass, - // then d2d-copies the residual into the device slot. No CacheGraphScope, - // no kCache*Name hook plumbing beyond the shared feature node name. ---- + // then d2d-copies the residual into the device slot. ---- sd::Tensor compute_substep_capture(int n_threads, const sd::Tensor& x, const sd::Tensor& timesteps, @@ -3651,55 +3434,75 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* return std::move(pass.output); } - sd::DiffusionCacheResult compute_probe(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context, - const std::vector>& ref_latents, - bool increase_ref_index, - int probe_depth, - const void* branch_key = nullptr) { - sd::CacheGraphScope scope; - scope.mode = sd::CacheGraphScope::Mode::Probe; - scope.probe_depth = probe_depth; - // GPU DiCache path: attach the persistent prev_* tensors so - // build_probe_metrics emits delta_y/delta_x/gamma reductions in-graph; - // only scalars are read back. - if (dicache_gpu_enabled() && branch_key != nullptr) { - auto it = dicache_gpu_states_.find(branch_key); - if (it != dicache_gpu_states_.end() && it->second.has_probe_hist) { - DiCacheGpuState& gpu = it->second; - scope.gpu_metric = true; - scope.prev_probe_node = gpu.prev_probe; - scope.prev_input_node = gpu.prev_input; - if (gpu.probe_resid_count >= 2) { - scope.probe_prev1_node = gpu.probe_prev1; - scope.probe_prev2_node = gpu.probe_prev2; - } - } + // ---- Substep-path tap-driven DiCache seed capture. A full + // forward whose post-readback refreshes the persistent DiCacheGpuState rings + // device-to-device from the tap-woven residual/probe nodes. No device slot is + // involved (DiCache's StateManager slot is host-backed/vestigial; the real + // cross-step state is the DiCacheGpuState ring). + void writeback_from_taps(const void* branch_key) { + // Shape source: the tap-woven residual (ModelOut - ModelIn). + // Allocate the rings lazily at that shape. + ggml_tensor* feat = get_cache_tensor_by_name("ed_cache_feature"); + if (feat == nullptr) { + return; } + DiCacheGpuState& s = + ensure_dicache_gpu_state(branch_key, feat->ne[0], feat->ne[1], feat->ne[2]); + if (s.buffer == nullptr) { + return; + } + const int m = std::max(1, dicache_probe_depth_); + const std::string probe_tap = "ed_tap:block_out[" + std::to_string(m - 1) + "]"; + // A silent copy failure (shape/name mismatch) would leave the ring stale + // and silently degrade the decision, so surface it loudly. + auto cp = [&](const char* name, ggml_tensor* dst) { + if (!copy_named_cache_tensor_to(name, dst)) { + LOG_ERROR("dicache writeback copy failed: %s", name); + } + }; + // full residual ring: prev2 <- prev1, prev1 <- feature + std::swap(s.resid_prev1, s.resid_prev2); + cp("ed_cache_feature", s.resid_prev1); + s.resid_count = std::min(2, s.resid_count + 1); + // probe residual ring: prev2 <- prev1, prev1 <- (probe - before) + std::swap(s.probe_prev1, s.probe_prev2); + cp("ed_cache_probe_resid", s.probe_prev1); + s.probe_resid_count = std::min(2, s.probe_resid_count + 1); + // raw probe & input for the delta_y / delta_x decision metrics + cp(probe_tap.c_str(), s.prev_probe); + cp("ed_tap:model_in", s.prev_input); + s.has_probe_hist = true; + } + + sd::Tensor compute_substep_capture_probe(int n_threads, + const sd::Tensor& x, + const sd::Tensor& timesteps, + const sd::Tensor& context, + const std::vector>& ref_latents, + bool increase_ref_index, + int probe_depth, + const void* branch_key) { + const int m = std::max(1, probe_depth); + edgedit::cache::TapRegistry reg; + // Full compute (no stop_after): block_out(m-1) is tapped in-loop anyway. + reg.set_requested({edgedit::cache::AnchorRef::model_in(), + edgedit::cache::AnchorRef::block_out(m - 1), + edgedit::cache::AnchorRef::model_out()}); + reg.set_capture_residual(true); // weave ed_cache_feature = ModelOut - ModelIn + reg.set_capture_writeback(m); // weave ed_cache_probe_resid = block_out(m-1) - ModelIn + std::function handoff = [&]() { writeback_from_taps(branch_key); }; auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, ref_latents, increase_ref_index); }; - auto pass = run_cache_pass(get_graph, n_threads, &scope, x.dim()); - sd::DiffusionCacheResult out; - if (scope.gpu_metric) { - // GPU path: no ~50MB host tensors; only scalars. - out.delta_y = pass.delta_y; - out.delta_x = pass.delta_x; - out.gamma = pass.gamma; - return out; - } - out.probe = pass.probe.empty() ? std::move(pass.output) : std::move(pass.probe); - out.before = std::move(pass.before); - return out; + auto pass = run_substep_pass(get_graph, n_threads, ®, x.dim(), {}, handoff); + return std::move(pass.output); } - // ---- Substep-path (ED_CACHE_SUBSTEP) tap-driven DiCache probe. Replaces the - // compute_probe GPU-metric path: requests ModelIn + BlockOut[m-1] taps, stops + // ---- Substep-path tap-driven DiCache probe. Requests + // ModelIn + BlockOut[m-1] taps, stops // the forward after m blocks, threads the persistent cross-step operands into // the registry, and the runner weaves delta_y/delta_x/gamma from the taps + - // operands. Reads back only the scalars. No CacheGraphScope. ---- + // operands. Reads back only the scalars. ---- sd::DiffusionCacheResult compute_substep_probe(int n_threads, const sd::Tensor& x, const sd::Tensor& timesteps, diff --git a/src/dit_models/models/wan.hpp b/src/dit_models/models/wan.hpp index 8588f70b..2aa255ff 100644 --- a/src/dit_models/models/wan.hpp +++ b/src/dit_models/models/wan.hpp @@ -4615,15 +4615,10 @@ namespace WAN { auto x_orig = x; // Cache seam: the block stack transforms the single stream `x`. - // Disabled under SP (sequence-sharded) and under VACE (the `c` stream - // is folded into `x` inside the loop; a skip would leave it stale). // The cached region is blocks [region_start, region_end); the default // whole-stack region matches the pre-region behaviour. - sd::CacheGraphScope* cache_scope = - (use_sp_mainline || c != nullptr) ? nullptr : ctx->cache_scope; - const bool cache_inject = cache_scope != nullptr && cache_scope->inject_mode(); // Substep tap: block-stack input anchor (ModelIn). Conditional no-op - // unless requested. Coexists with the legacy cache_scope path. + // unless requested. tap(ctx, edgedit::cache::AnchorRef::model_in(), x_orig); ggml_tensor* sp_prepared_pe = nullptr; if (use_sp_mainline) { @@ -4635,23 +4630,13 @@ namespace WAN { for (int i = 0; i < params.num_layers; i++) { // Tap-driven inject (substep reuse): at the region start, replace the - // stream with x_before + inject_input and jump past the region — no - // CacheGraphScope. x_orig is the block-stack input (ModelIn). + // stream with x_before + inject_input and jump past the region. + // x_orig is the block-stack input (ModelIn). if (ctx->tap_registry != nullptr && ctx->tap_registry->inject_at(i)) { x = build_tap_inject(ctx, x_orig); i = ctx->tap_registry->inject_resume() - 1; continue; } - if (cache_inject) { - if (ggml_tensor* injected = cache_scope->step_inject_region(ctx->ggml_ctx, i, x)) { - x = injected; - i = cache_scope->inject_resume_index(params.num_layers) - 1; - continue; - } - } - if (cache_scope != nullptr) { - cache_scope->begin_region(i, x); - } auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); @@ -4698,13 +4683,6 @@ namespace WAN { if (ctx->tap_registry != nullptr && ctx->tap_registry->stop_after(i)) { return x; } - if (cache_scope != nullptr) { - cache_scope->end_region(ctx->ggml_ctx, i, params.num_layers, x); - if (cache_scope->stop_after_block(i)) { - cache_scope->on_probe(x); - return x; - } - } } // Substep tap: block-stack output anchor (ModelOut) — the residual's @@ -5069,75 +5047,9 @@ namespace WAN { return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); } - // ---- Cache seam passes (Feature/Probe policies). VACE never active on - // these passes (the model gates the seam off when a VACE stream exists). ---- - sd::DiffusionCacheResult compute_capture(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context, - const sd::Tensor& clip_fea, - const sd::Tensor& c_concat, - int region_start = 0, - int region_end = -1) { - sd::CacheGraphScope scope; - scope.mode = sd::CacheGraphScope::Mode::Capture; - scope.region_start = region_start; - scope.region_end = region_end; - auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(x, timesteps, context, clip_fea, c_concat, {}, {}, 1.f); - }; - auto pass = run_cache_pass(get_graph, n_threads, &scope, x.dim()); - sd::DiffusionCacheResult out; - out.output = std::move(pass.output); - out.feature = std::move(pass.feature); - return out; - } - - sd::Tensor compute_inject(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context, - const sd::Tensor& clip_fea, - const sd::Tensor& c_concat, - const sd::Tensor& feature, - int region_start = 0, - int region_end = -1) { - sd::CacheGraphScope scope; - scope.mode = sd::CacheGraphScope::Mode::Inject; - scope.region_start = region_start; - scope.region_end = region_end; - inject_feature_host_ = feature; - auto get_graph = [&]() -> ggml_cgraph* { - scope.inject_feature = make_input(inject_feature_host_); - return build_graph(x, timesteps, context, clip_fea, c_concat, {}, {}, 1.f); - }; - auto pass = run_cache_pass(get_graph, n_threads, &scope, x.dim()); - return std::move(pass.output); - } - - sd::DiffusionCacheResult compute_probe(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context, - const sd::Tensor& clip_fea, - const sd::Tensor& c_concat, - int probe_depth) { - sd::CacheGraphScope scope; - scope.mode = sd::CacheGraphScope::Mode::Probe; - scope.probe_depth = probe_depth; - auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(x, timesteps, context, clip_fea, c_concat, {}, {}, 1.f); - }; - auto pass = run_cache_pass(get_graph, n_threads, &scope, x.dim()); - sd::DiffusionCacheResult out; - out.probe = pass.probe.empty() ? std::move(pass.output) : std::move(pass.probe); - out.before = std::move(pass.before); - return out; - } - - // ---- Substep-path (ED_CACHE_SUBSTEP) tap-driven capture (host path). Wan + // ---- Substep-path tap-driven capture (host path). Wan // has no device slot, so the residual is read back to host (feature). The - // runner weaves (ModelOut - ModelIn) from the taps; no CacheGraphScope. ---- + // runner weaves (ModelOut - ModelIn) from the taps. ---- sd::DiffusionCacheResult compute_substep_capture(int n_threads, const sd::Tensor& x, const sd::Tensor& timesteps, @@ -5159,9 +5071,9 @@ namespace WAN { return out; } - // ---- Substep-path (ED_CACHE_SUBSTEP) tap-driven probe (host path). Requests + // ---- Substep-path tap-driven probe (host path). Requests // ModelIn + BlockOut[m-1] taps, stops after m blocks, reads the before/probe - // tensors back to host for the host DiCache metric. No CacheGraphScope. ---- + // tensors back to host for the host DiCache metric. ---- sd::DiffusionCacheResult compute_substep_probe(int n_threads, const sd::Tensor& x, const sd::Tensor& timesteps, @@ -5188,10 +5100,10 @@ namespace WAN { return out; } - // ---- Substep-path (ED_CACHE_SUBSTEP) tap-driven inject (host reuse). Uploads + // ---- Substep-path tap-driven inject (host reuse). Uploads // the host residual as a graph input and drives the forward's registry inject: // at the region start the stream becomes x_before + inject_input and the - // region's blocks are skipped. No CacheGraphScope. ---- + // region's blocks are skipped. ---- sd::Tensor compute_substep_inject(int n_threads, const sd::Tensor& x, const sd::Tensor& timesteps, diff --git a/src/dit_models/pipelines/flux_kontext_pipeline.cpp b/src/dit_models/pipelines/flux_kontext_pipeline.cpp index 6cd59f38..7aeff305 100644 --- a/src/dit_models/pipelines/flux_kontext_pipeline.cpp +++ b/src/dit_models/pipelines/flux_kontext_pipeline.cpp @@ -1015,29 +1015,6 @@ bool FluxKontextPipeline::generate_one_image(const ed_image_generation_params_t* cond_in.c_crossattn, {}, cond_in.c_vector, guidance, ref_latents); }; - const bool seam_ok = !use_cfg_parallel && flux_runner_->feature_cache_available(); - hooks.feature_supported = seam_ok; - if (seam_ok) { - hooks.capture = [&](int region_start, int region_end) { - return flux_runner_->compute_capture(n_threads, noised_input, timesteps, - cond_in.c_crossattn, {}, cond_in.c_vector, - guidance, ref_latents, false, - region_start, region_end); - }; - hooks.inject = [&](const sd::Tensor& feat, int region_start, int region_end) { - return flux_runner_->compute_inject(n_threads, noised_input, timesteps, - cond_in.c_crossattn, {}, cond_in.c_vector, - guidance, ref_latents, false, feat, - region_start, region_end); - }; - if (cache_runtime.granularity() == cache::CacheGranularity::Probe) { - hooks.probe = [&](int depth) { - return flux_runner_->compute_probe(n_threads, noised_input, timesteps, - cond_in.c_crossattn, {}, cond_in.c_vector, - guidance, ref_latents, false, depth); - }; - } - } return hooks; }; diff --git a/src/dit_models/pipelines/flux_pipeline.cpp b/src/dit_models/pipelines/flux_pipeline.cpp index 06d2e574..884ce8e7 100644 --- a/src/dit_models/pipelines/flux_pipeline.cpp +++ b/src/dit_models/pipelines/flux_pipeline.cpp @@ -933,7 +933,6 @@ bool FluxPipeline::generate_one_image(const ed_image_generation_params_t* params cond_in.c_crossattn, {}, cond_in.c_vector, guidance); }; const bool seam_ok = !use_cfg_parallel && flux_runner_->feature_cache_available(); - hooks.feature_supported = seam_ok; if (seam_ok) { const void* branch_key = static_cast(&cond_in); // Only DiCache (Probe) uses the on-GPU probe/inject seam that a @@ -947,31 +946,14 @@ bool FluxPipeline::generate_one_image(const ed_image_generation_params_t* params const bool feature_gpu = !is_probe && cache_runtime.granularity() == cache::CacheGranularity::Feature && Flux::FluxRunner::feature_gpu_enabled(); - const void* capture_key = is_probe ? branch_key : nullptr; if (feature_gpu) { - // Declarative device-slot seam (B2): the lowering hands us the - // slot's device tensor; capture stores the residual into it and - // reuse injects x_before + slot on-device (no host round-trip). - hooks.capture_to_slot = [&](const std::function&)>& alloc_slot, - int region_start, int region_end) { - return flux_runner_->compute_capture_to_slot( - n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, - cond_in.c_vector, guidance, {}, false, - alloc_slot, region_start, region_end); - }; - // Substep-path tap-driven capture (ED_CACHE_SUBSTEP): TapRegistry, + // Substep-path tap-driven capture: TapRegistry, // not CacheGraphScope. hooks.substep_capture = [&](const std::function&)>& alloc_slot) { return flux_runner_->compute_substep_capture( n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, cond_in.c_vector, guidance, {}, false, alloc_slot); }; - hooks.inject_from_slot = [&](void* slot, int region_start, int region_end) { - return flux_runner_->compute_inject_from_slot( - n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, - cond_in.c_vector, guidance, {}, false, - static_cast(slot), region_start, region_end); - }; // Substep-path tap-driven device inject (MagCache): registry inject. hooks.substep_inject_slot = [&](void* slot, int region_start, int region_end) { return flux_runner_->compute_substep_inject_slot( @@ -979,26 +961,9 @@ bool FluxPipeline::generate_one_image(const ed_image_generation_params_t* params cond_in.c_vector, guidance, {}, false, static_cast(slot), region_start, region_end); }; - } else { - hooks.capture = [&, capture_key](int region_start, int region_end) { - return flux_runner_->compute_capture(n_threads, noised_input, timesteps, - cond_in.c_crossattn, {}, cond_in.c_vector, - guidance, {}, false, region_start, region_end, - capture_key); - }; } - hooks.inject = [&](const sd::Tensor& feat, int region_start, int region_end) { - return flux_runner_->compute_inject(n_threads, noised_input, timesteps, - cond_in.c_crossattn, {}, cond_in.c_vector, - guidance, {}, false, feat, region_start, region_end); - }; if (cache_runtime.granularity() == cache::CacheGranularity::Probe) { - hooks.probe = [&, branch_key](int depth) { - return flux_runner_->compute_probe(n_threads, noised_input, timesteps, - cond_in.c_crossattn, {}, cond_in.c_vector, - guidance, {}, false, depth, branch_key); - }; - // Substep-path tap-driven probe (ED_CACHE_SUBSTEP): delta_y/gamma + // Substep-path tap-driven probe: delta_y/gamma // on-device from taps + persistent operands, no CacheGraphScope. const bool delta_minus = cache_runtime.dicache_delta_minus(); hooks.substep_probe = [&, branch_key, delta_minus](int depth) { @@ -1012,12 +977,6 @@ bool FluxPipeline::generate_one_image(const ed_image_generation_params_t* params // the lowering take the declarative host probe path (the residual // ring blend), instead of the legacy on-device reconstruction. if (Flux::FluxRunner::dicache_gpu_enabled()) { - hooks.inject_gpu = [&, branch_key](float gamma, int region_start, int region_end) { - return flux_runner_->compute_inject_gpu(n_threads, noised_input, timesteps, - cond_in.c_crossattn, {}, cond_in.c_vector, - guidance, {}, false, gamma, branch_key, - region_start, region_end); - }; // Substep-path tap-driven device inject (DiCache gamma-blend). hooks.substep_inject_gpu = [&, branch_key](float gamma, int region_start, int region_end) { return flux_runner_->compute_substep_inject_gpu(n_threads, noised_input, timesteps, @@ -1025,6 +984,15 @@ bool FluxPipeline::generate_one_image(const ed_image_generation_params_t* params guidance, {}, false, gamma, branch_key, region_start, region_end); }; + // Substep-path tap-driven seed capture: full forward that + // refreshes the DiCacheGpuState rings device-to-device (replaces + // the legacy compute_capture / run_cache_pass path). + const int probe_depth = cache_runtime.dicache_probe_depth(); + hooks.substep_capture_probe = [&, branch_key, probe_depth]() { + return flux_runner_->compute_substep_capture_probe( + n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, cond_in.c_vector, + guidance, {}, false, probe_depth, branch_key); + }; } } } diff --git a/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp b/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp index 69b45a5a..e2ea486c 100644 --- a/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp +++ b/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp @@ -762,65 +762,6 @@ bool QwenImageEditPipeline::generate_one_image(const ed_image_generation_params_ return diffusion_->compute(n_threads, x, timesteps, cond_in.c_crossattn, ref_latents, true); }; - const bool seam_ok = !use_cfg_parallel && diffusion_->feature_cache_available(); - hooks.feature_supported = seam_ok; - if (seam_ok) { - const void* branch_key = static_cast(&cond_in); - // Only DiCache (Probe) uses the on-GPU probe/inject seam that a - // branch_key drives; passing it into compute_capture for a Feature - // method (MagCache/TaylorSeer/SenCache) would flip gpu_metric on and - // suppress the host feature readback those methods rely on. - const bool is_probe = cache_runtime.granularity() == cache::CacheGranularity::Probe; - // Feature-granularity on-GPU reuse: keep the captured residual on - // device and inject it there on skips, avoiding the ~50MB host - // reconstruct copy + H2D upload the host inject path pays per skip. - const bool feature_gpu = !is_probe && - cache_runtime.granularity() == cache::CacheGranularity::Feature && - Qwen::QwenImageRunner::feature_gpu_enabled(); - const void* capture_key = is_probe ? branch_key : nullptr; - if (feature_gpu) { - // Declarative device-slot seam (B2): the lowering hands us the - // slot's device tensor; capture stores the residual into it and - // reuse injects x_before + slot on-device (no host round-trip). - hooks.capture_to_slot = [&, cond_in](const std::function&)>& alloc_slot, - int region_start, int region_end) { - return diffusion_->compute_capture_to_slot( - n_threads, x, timesteps, cond_in.c_crossattn, ref_latents, true, - alloc_slot, region_start, region_end); - }; - hooks.inject_from_slot = [&, cond_in](void* slot, int region_start, int region_end) { - return diffusion_->compute_inject_from_slot( - n_threads, x, timesteps, cond_in.c_crossattn, ref_latents, true, - static_cast(slot), region_start, region_end); - }; - } else { - hooks.capture = [&, cond_in, capture_key](int region_start, int region_end) { - return diffusion_->compute_capture(n_threads, x, timesteps, cond_in.c_crossattn, - ref_latents, true, region_start, region_end, - capture_key); - }; - } - hooks.inject = [&, cond_in](const sd::Tensor& feat, int region_start, int region_end) { - return diffusion_->compute_inject(n_threads, x, timesteps, cond_in.c_crossattn, - ref_latents, true, feat, region_start, region_end); - }; - if (cache_runtime.granularity() == cache::CacheGranularity::Probe) { - hooks.probe = [&, cond_in, branch_key](int depth) { - return diffusion_->compute_probe(n_threads, x, timesteps, cond_in.c_crossattn, - ref_latents, true, depth, branch_key); - }; - // Only wire on-GPU inject when the model's GPU DiCache path is - // active; with ED_DICACHE_GPU=0 the lowering takes the declarative - // host probe path instead (residual-ring blend). - if (Qwen::QwenImageRunner::dicache_gpu_enabled()) { - hooks.inject_gpu = [&, cond_in, branch_key](float gamma, int region_start, int region_end) { - return diffusion_->compute_inject_gpu(n_threads, x, timesteps, cond_in.c_crossattn, - ref_latents, true, gamma, branch_key, - region_start, region_end); - }; - } - } - } return hooks; }; diff --git a/src/dit_models/pipelines/qwen_image_pipeline.cpp b/src/dit_models/pipelines/qwen_image_pipeline.cpp index 0eeb8f69..8233797c 100644 --- a/src/dit_models/pipelines/qwen_image_pipeline.cpp +++ b/src/dit_models/pipelines/qwen_image_pipeline.cpp @@ -699,7 +699,6 @@ bool QwenImagePipeline::generate_one_image(const ed_image_generation_params_t* p empty_ref_latents, false); }; const bool seam_ok = !use_cfg_parallel && diffusion_->feature_cache_available(); - hooks.feature_supported = seam_ok; if (seam_ok) { const void* branch_key = static_cast(&cond_in); // Only DiCache (Probe) uses the on-GPU probe/inject seam that a @@ -713,29 +712,14 @@ bool QwenImagePipeline::generate_one_image(const ed_image_generation_params_t* p const bool feature_gpu = !is_probe && cache_runtime.granularity() == cache::CacheGranularity::Feature && Qwen::QwenImageRunner::feature_gpu_enabled(); - const void* capture_key = is_probe ? branch_key : nullptr; if (feature_gpu) { - // Declarative device-slot seam (B2): the lowering hands us the - // slot's device tensor; capture stores the residual into it and - // reuse injects x_before + slot on-device (no host round-trip). - hooks.capture_to_slot = [&, cond_in](const std::function&)>& alloc_slot, - int region_start, int region_end) { - return diffusion_->compute_capture_to_slot( - n_threads, x, timesteps, cond_in.c_crossattn, empty_ref_latents, false, - alloc_slot, region_start, region_end); - }; - // Substep-path tap-driven capture (ED_CACHE_SUBSTEP): same slot + // Substep-path tap-driven capture: same slot // contract, but driven through the TapRegistry, not CacheGraphScope. hooks.substep_capture = [&, cond_in](const std::function&)>& alloc_slot) { return diffusion_->compute_substep_capture( n_threads, x, timesteps, cond_in.c_crossattn, empty_ref_latents, false, alloc_slot); }; - hooks.inject_from_slot = [&, cond_in](void* slot, int region_start, int region_end) { - return diffusion_->compute_inject_from_slot( - n_threads, x, timesteps, cond_in.c_crossattn, empty_ref_latents, false, - static_cast(slot), region_start, region_end); - }; // Substep-path tap-driven device inject (MagCache): x_before + slot // via the forward's registry inject, no CacheGraphScope. hooks.substep_inject_slot = [&, cond_in](void* slot, int region_start, int region_end) { @@ -743,23 +727,9 @@ bool QwenImagePipeline::generate_one_image(const ed_image_generation_params_t* p n_threads, x, timesteps, cond_in.c_crossattn, empty_ref_latents, false, static_cast(slot), region_start, region_end); }; - } else { - hooks.capture = [&, cond_in, capture_key](int region_start, int region_end) { - return diffusion_->compute_capture(n_threads, x, timesteps, cond_in.c_crossattn, - empty_ref_latents, false, region_start, region_end, - capture_key); - }; } - hooks.inject = [&, cond_in](const sd::Tensor& feat, int region_start, int region_end) { - return diffusion_->compute_inject(n_threads, x, timesteps, cond_in.c_crossattn, - empty_ref_latents, false, feat, region_start, region_end); - }; if (cache_runtime.granularity() == cache::CacheGranularity::Probe) { - hooks.probe = [&, cond_in, branch_key](int depth) { - return diffusion_->compute_probe(n_threads, x, timesteps, cond_in.c_crossattn, - empty_ref_latents, false, depth, branch_key); - }; - // Substep-path tap-driven probe (ED_CACHE_SUBSTEP): delta_y/gamma + // Substep-path tap-driven probe: delta_y/gamma // computed on-device from taps + persistent operands, no scope. const bool delta_minus = cache_runtime.dicache_delta_minus(); hooks.substep_probe = [&, cond_in, branch_key, delta_minus](int depth) { @@ -771,17 +741,21 @@ bool QwenImagePipeline::generate_one_image(const ed_image_generation_params_t* p // active; with ED_DICACHE_GPU=0 the lowering takes the declarative // host probe path instead (residual-ring blend). if (Qwen::QwenImageRunner::dicache_gpu_enabled()) { - hooks.inject_gpu = [&, cond_in, branch_key](float gamma, int region_start, int region_end) { - return diffusion_->compute_inject_gpu(n_threads, x, timesteps, cond_in.c_crossattn, - empty_ref_latents, false, gamma, branch_key, - region_start, region_end); - }; // Substep-path tap-driven device inject (DiCache gamma-blend). hooks.substep_inject_gpu = [&, cond_in, branch_key](float gamma, int region_start, int region_end) { return diffusion_->compute_substep_inject_gpu(n_threads, x, timesteps, cond_in.c_crossattn, empty_ref_latents, false, gamma, branch_key, region_start, region_end); }; + // Substep-path tap-driven seed capture: full forward that + // refreshes the DiCacheGpuState rings device-to-device (replaces + // the legacy compute_capture / run_cache_pass path). + const int probe_depth = cache_runtime.dicache_probe_depth(); + hooks.substep_capture_probe = [&, cond_in, branch_key, probe_depth]() { + return diffusion_->compute_substep_capture_probe( + n_threads, x, timesteps, cond_in.c_crossattn, empty_ref_latents, + false, probe_depth, branch_key); + }; } } } diff --git a/src/dit_models/pipelines/sd3_pipeline.cpp b/src/dit_models/pipelines/sd3_pipeline.cpp index 153e77c4..ac19e05b 100644 --- a/src/dit_models/pipelines/sd3_pipeline.cpp +++ b/src/dit_models/pipelines/sd3_pipeline.cpp @@ -467,15 +467,8 @@ bool SD3Pipeline::generate_one_image(const ed_image_generation_params_t* params, return diffusion_->compute(n_threads, p); }; const bool seam_ok = !use_cfg_parallel && diffusion_->supports_feature_cache(); - hooks.feature_supported = seam_ok; if (seam_ok) { - hooks.capture = [&, cond_in](int region_start, int region_end) { - DiffusionParams p = diffusion_params; - p.context = &cond_in.c_crossattn; - p.y = &cond_in.c_vector; - return diffusion_->compute_capture(n_threads, p, region_start, region_end); - }; - // Substep-path tap-driven host capture (ED_CACHE_SUBSTEP): residual + // Substep-path tap-driven host capture: residual // via ModelIn/ModelOut taps, read back to host. No CacheGraphScope. hooks.substep_capture_host = [&, cond_in]() { DiffusionParams p = diffusion_params; @@ -483,13 +476,7 @@ bool SD3Pipeline::generate_one_image(const ed_image_generation_params_t* params, p.y = &cond_in.c_vector; return diffusion_->compute_substep_capture_host(n_threads, p); }; - hooks.inject = [&, cond_in](const sd::Tensor& feat, int region_start, int region_end) { - DiffusionParams p = diffusion_params; - p.context = &cond_in.c_crossattn; - p.y = &cond_in.c_vector; - return diffusion_->compute_inject(n_threads, p, feat, region_start, region_end); - }; - // Substep-path tap-driven host inject (ED_CACHE_SUBSTEP). + // Substep-path tap-driven host inject. hooks.substep_inject_host = [&, cond_in](const sd::Tensor& feat) { DiffusionParams p = diffusion_params; p.context = &cond_in.c_crossattn; @@ -497,13 +484,7 @@ bool SD3Pipeline::generate_one_image(const ed_image_generation_params_t* params, return diffusion_->compute_substep_inject_host(n_threads, p, feat, 0, -1); }; if (cache_runtime.granularity() == cache::CacheGranularity::Probe) { - hooks.probe = [&, cond_in](int depth) { - DiffusionParams p = diffusion_params; - p.context = &cond_in.c_crossattn; - p.y = &cond_in.c_vector; - return diffusion_->compute_probe(n_threads, p, depth); - }; - // Substep-path tap-driven host probe (ED_CACHE_SUBSTEP). + // Substep-path tap-driven host probe. hooks.substep_probe_host = [&, cond_in](int depth) { DiffusionParams p = diffusion_params; p.context = &cond_in.c_crossattn; diff --git a/src/dit_models/pipelines/wan_pipeline.cpp b/src/dit_models/pipelines/wan_pipeline.cpp index 7f20ebf3..3847bd47 100644 --- a/src/dit_models/pipelines/wan_pipeline.cpp +++ b/src/dit_models/pipelines/wan_pipeline.cpp @@ -1040,7 +1040,6 @@ sd::Tensor WanPipeline::euler_denoise(const std::shared_ptrsupports_feature_cache() && model->feature_cache_available(); - hooks.feature_supported = seam_ok; if (seam_ok) { auto set_params = [&, cond_in]() { diffusion_params.context = cond_in.c_crossattn.empty() ? nullptr : &cond_in.c_crossattn; @@ -1050,32 +1049,20 @@ sd::Tensor WanPipeline::euler_denoise(const std::shared_ptrcompute_capture(runtime_->n_threads(), diffusion_params, region_start, region_end); - }; - // Substep-path tap-driven host capture (ED_CACHE_SUBSTEP): residual + // Substep-path tap-driven host capture: residual // via ModelIn/ModelOut taps, read back to host. No CacheGraphScope. hooks.substep_capture_host = [&, set_params]() { set_params(); return model->compute_substep_capture_host(runtime_->n_threads(), diffusion_params); }; - hooks.inject = [&, set_params](const sd::Tensor& feat, int region_start, int region_end) { - set_params(); - return model->compute_inject(runtime_->n_threads(), diffusion_params, feat, region_start, region_end); - }; - // Substep-path tap-driven host inject (ED_CACHE_SUBSTEP): x_before + + // Substep-path tap-driven host inject: x_before + // feature with the region skipped, no CacheGraphScope. hooks.substep_inject_host = [&, set_params](const sd::Tensor& feat) { set_params(); return model->compute_substep_inject_host(runtime_->n_threads(), diffusion_params, feat, 0, -1); }; if (cache_runtime.granularity() == cache::CacheGranularity::Probe) { - hooks.probe = [&, set_params](int depth) { - set_params(); - return model->compute_probe(runtime_->n_threads(), diffusion_params, depth); - }; - // Substep-path tap-driven host probe (ED_CACHE_SUBSTEP): before/probe + // Substep-path tap-driven host probe: before/probe // via ModelIn/BlockOut[m-1] taps, read back to host. hooks.substep_probe_host = [&, set_params](int depth) { set_params(); From 416e5cfe881f1ecce10d8732e2a0bf52a5d6917e Mon Sep 17 00:00:00 2001 From: Liwei Lan Date: Wed, 15 Jul 2026 17:24:47 +0800 Subject: [PATCH 2/8] fix: wire substep cache hooks in flux_kontext & qwen_image_edit variants --- .../pipelines/flux_kontext_pipeline.cpp | 61 ++++++++++++++++++- .../pipelines/qwen_image_edit_pipeline.cpp | 41 +++++++++++++ 2 files changed, 101 insertions(+), 1 deletion(-) diff --git a/src/dit_models/pipelines/flux_kontext_pipeline.cpp b/src/dit_models/pipelines/flux_kontext_pipeline.cpp index 7aeff305..ab9fdcee 100644 --- a/src/dit_models/pipelines/flux_kontext_pipeline.cpp +++ b/src/dit_models/pipelines/flux_kontext_pipeline.cpp @@ -963,9 +963,24 @@ bool FluxKontextPipeline::generate_one_image(const ed_image_generation_params_t* parallel::cfg_parallel_available(runtime_->parallel_context()); const bool cache_seam_available = !cache_use_cfg_parallel && flux_runner_->feature_cache_available(); + // Wire the device store only when the on-GPU feature-reuse path is active + // (ED_FEATURE_CACHE_GPU on); with it off, leave the store null so a + // device_backed slot cleanly falls back to the host declarative path. + cache::ICacheDeviceStore* cache_store = + (cache_seam_available && flux_runner_ != nullptr && + Flux::FluxRunner::feature_gpu_enabled()) + ? flux_runner_->cache_device_store() + : nullptr; const bool cache_enabled = - cache_runtime.init(params->sample, version_, sigmas, cache_seam_available, nullptr, + cache_runtime.init(params->sample, version_, sigmas, cache_seam_available, cache_store, cache_use_cfg_parallel); + // GPU DiCache (ED_DICACHE_GPU): reset per-generation persistent state and set + // the probe depth the capture step uses to snapshot its probe residual. Read + // the resolved depth from the engine so it stays in sync with the policy config. + if (cache_enabled && flux_runner_ != nullptr) { + flux_runner_->reset_dicache_gpu_states(); + flux_runner_->dicache_probe_depth_ = cache_runtime.dicache_probe_depth(); + } const int64_t sample_start_ms = ggml_time_ms(); GenerationControl* control = runtime_ != nullptr ? runtime_->generation_control() : nullptr; for (int step = 0; step < steps; ++step) { @@ -1015,6 +1030,50 @@ bool FluxKontextPipeline::generate_one_image(const ed_image_generation_params_t* cond_in.c_crossattn, {}, cond_in.c_vector, guidance, ref_latents); }; + const bool seam_ok = !use_cfg_parallel && flux_runner_->feature_cache_available(); + if (seam_ok) { + const void* branch_key = static_cast(&cond_in); + const bool is_probe = cache_runtime.granularity() == cache::CacheGranularity::Probe; + const bool feature_gpu = !is_probe && + cache_runtime.granularity() == cache::CacheGranularity::Feature && + Flux::FluxRunner::feature_gpu_enabled(); + if (feature_gpu) { + hooks.substep_capture = [&](const std::function&)>& alloc_slot) { + return flux_runner_->compute_substep_capture( + n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, + cond_in.c_vector, guidance, ref_latents, false, alloc_slot); + }; + hooks.substep_inject_slot = [&](void* slot, int region_start, int region_end) { + return flux_runner_->compute_substep_inject_slot( + n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, + cond_in.c_vector, guidance, ref_latents, false, + static_cast(slot), region_start, region_end); + }; + } + if (cache_runtime.granularity() == cache::CacheGranularity::Probe) { + const bool delta_minus = cache_runtime.dicache_delta_minus(); + hooks.substep_probe = [&, branch_key, delta_minus](int depth) { + return flux_runner_->compute_substep_probe(n_threads, noised_input, timesteps, + cond_in.c_crossattn, {}, cond_in.c_vector, + guidance, ref_latents, false, depth, branch_key, + delta_minus); + }; + if (Flux::FluxRunner::dicache_gpu_enabled()) { + hooks.substep_inject_gpu = [&, branch_key](float gamma, int region_start, int region_end) { + return flux_runner_->compute_substep_inject_gpu(n_threads, noised_input, timesteps, + cond_in.c_crossattn, {}, cond_in.c_vector, + guidance, ref_latents, false, gamma, branch_key, + region_start, region_end); + }; + const int probe_depth = cache_runtime.dicache_probe_depth(); + hooks.substep_capture_probe = [&, branch_key, probe_depth]() { + return flux_runner_->compute_substep_capture_probe( + n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, cond_in.c_vector, + guidance, ref_latents, false, probe_depth, branch_key); + }; + } + } + } return hooks; }; diff --git a/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp b/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp index e2ea486c..2a3161b2 100644 --- a/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp +++ b/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp @@ -762,6 +762,47 @@ bool QwenImageEditPipeline::generate_one_image(const ed_image_generation_params_ return diffusion_->compute(n_threads, x, timesteps, cond_in.c_crossattn, ref_latents, true); }; + const bool seam_ok = !use_cfg_parallel && diffusion_->feature_cache_available(); + if (seam_ok) { + const void* branch_key = static_cast(&cond_in); + const bool is_probe = cache_runtime.granularity() == cache::CacheGranularity::Probe; + const bool feature_gpu = !is_probe && + cache_runtime.granularity() == cache::CacheGranularity::Feature && + Qwen::QwenImageRunner::feature_gpu_enabled(); + if (feature_gpu) { + hooks.substep_capture = [&, cond_in](const std::function&)>& alloc_slot) { + return diffusion_->compute_substep_capture( + n_threads, x, timesteps, cond_in.c_crossattn, ref_latents, true, + alloc_slot); + }; + hooks.substep_inject_slot = [&, cond_in](void* slot, int region_start, int region_end) { + return diffusion_->compute_substep_inject_slot( + n_threads, x, timesteps, cond_in.c_crossattn, ref_latents, true, + static_cast(slot), region_start, region_end); + }; + } + if (cache_runtime.granularity() == cache::CacheGranularity::Probe) { + const bool delta_minus = cache_runtime.dicache_delta_minus(); + hooks.substep_probe = [&, cond_in, branch_key, delta_minus](int depth) { + return diffusion_->compute_substep_probe(n_threads, x, timesteps, + cond_in.c_crossattn, ref_latents, + true, depth, branch_key, delta_minus); + }; + if (Qwen::QwenImageRunner::dicache_gpu_enabled()) { + hooks.substep_inject_gpu = [&, cond_in, branch_key](float gamma, int region_start, int region_end) { + return diffusion_->compute_substep_inject_gpu(n_threads, x, timesteps, cond_in.c_crossattn, + ref_latents, true, gamma, branch_key, + region_start, region_end); + }; + const int probe_depth = cache_runtime.dicache_probe_depth(); + hooks.substep_capture_probe = [&, cond_in, branch_key, probe_depth]() { + return diffusion_->compute_substep_capture_probe( + n_threads, x, timesteps, cond_in.c_crossattn, ref_latents, + true, probe_depth, branch_key); + }; + } + } + } return hooks; }; From 070fafc50d87444019d9ea3556bb556cee1b37fc Mon Sep 17 00:00:00 2001 From: Liwei Lan Date: Wed, 15 Jul 2026 17:25:20 +0800 Subject: [PATCH 3/8] docs: refresh cache framework status to reflect completed dehook --- docs/cache_framework_status_and_roadmap.md | 263 +++++++++++++-------- 1 file changed, 164 insertions(+), 99 deletions(-) diff --git a/docs/cache_framework_status_and_roadmap.md b/docs/cache_framework_status_and_roadmap.md index 37533e5e..887a1ae4 100644 --- a/docs/cache_framework_status_and_roadmap.md +++ b/docs/cache_framework_status_and_roadmap.md @@ -12,7 +12,7 @@ ## 1. 背景与范围 `refactor/cache_framework` 把 step-caching 从旧的 `CacheController` 重构为声明式的 -`CacheEngine`（相对 `main`：89 文件、+14k/-4.2k 行）。本文做两件事： +`CacheEngine`。本文做两件事： 1. **盘点现状**：哪些能力真正落地、优点与差距分别是什么； 2. **给出前进目标**：把「落地声明式引擎」作为推荐方向深入分析其优势与问题，并列出备选方向。 @@ -20,6 +20,14 @@ 本文以 `docs/cache_framework_redesign.md` 描述的声明式架构为**目标态**，衡量「已落地 vs 差距」。文中所有「未承重 / 死代码」类断言均已由代码检索核实，并附 `file:line`。 +> **最新进展（2026-07-15）：cache 层与模型层的解耦已完成。** 07-14 之后的一批提交 +> 把「去 hook（dehook）」这一原本列为「数周量级、终态」的工作做完了：模型 forward 不再 +> 认识任何 cache 概念，改为在结构地标处调用 `TapRegistry` 的条件 `tap()`；旧的 +> `CacheGraphScope` seam 已被**整块删除**（commit `cd0c9b7`，`-1667` 行），`ED_CACHE_SUBSTEP` +> 门控随之消失——**substep 循环成为唯一执行路径**。下方 §2.6、§4、§5 已据此改写； +> §5 的三层路线图中，前两层（StateManager / Operator）＋原「Slice 4 去 hook」均已落地， +> 只剩第 3 层（GraphRepository）与 SP+缓存复合未做。 + > **术语：承重 vs 脚手架。**「承重」指该组件在真实推理路径上被调用、去掉会改变行为； > 「脚手架」指已实现但当前无任何调用方、去掉不影响任何现有功能的代码。本文的核心结论 > 之一，就是区分声明式设计里哪些是承重、哪些是脚手架。 @@ -52,11 +60,13 @@ vs 差距」。文中所有「未承重 / 死代码」类断言均已由代码 ### 2.4 三种 granularity（承重） -在 `cache_graph_lowering.cpp` 的 `execute` 中按 granularity 分流： +在 `cache_graph_lowering.cpp` 的 `execute_substeps` 中按 `SubstepOpKind` 分流（policy 的 +`next_substep()` 逐子步产出 `SubstepPlan`，中间层翻译为 tap 驱动的 runner pass）： - **Output**：黑盒，只用 `hooks.full()` + 主机侧输出 diff/复用，SP 下也可用； -- **Feature**：用 block-stack seam 捕获/注入残差（MagCache、TaylorSeer、SenCache）； -- **Probe**：跑浅层前缀再决策（DiCache）。 +- **Feature**：经 `TapRegistry` 在 ModelIn/ModelOut 锚点捕获/注入残差（MagCache、TaylorSeer、 + SenCache）； +- **Probe**：跑浅层前缀（`stop_after`）再决策（DiCache）。 ### 2.5 性能优化真实有效 @@ -66,13 +76,43 @@ vs 差距」。文中所有「未承重 / 死代码」类断言均已由代码 build-once 图复用； - skip 成本从 ~40ms/skip 降到 ~1ms，且验证 byte-identical，质量不变； - Flux：MagCache 2.05x、DiCache 默认 1.24x（近无损，PSNR ~40）。 +- **2026-07-14 更新**：Flux 和 Qwen 的默认 on-GPU 特征复用现统一走**声明式设备槽** + （`CacheStateManager` 设备后端），已替换 legacy `DiCacheGpuState` 路径；设备槽 vs + host-declarative 验证**字节一致（PSNR 100 dB）**，skip 计数逐位一致（见 §5.5）。 + 注：本轮 Qwen 验证需 `ED_QWEN_SINGLE_FUSED_ATTENTION=0`（默认融合注意力在当前 CUDA-13 + 环境下产出纯白图，属既有环境问题，与 cache 无关）。 + +### 2.6 cache/模型解耦已完成：TapRegistry 取代 seam（承重，2026-07-15） + +原设计文档把「彻底去 hook」列为终态、数周量级的 Slice 4。**它已经做完了**： + +- **依赖方向已倒置干净**：cache 层（`src/core/optimization/cache/`）**零 include 任何具体 + 模型 header**（flux/qwen/wan/mmdit），只依赖模型侧暴露的抽象——`anchor` / `cache_site` / + `model_topology` / `model_schema` / `tap_registry` / `model_cache_contract`。已由 grep 核实。 +- **模型 forward 不再认识 cache**：4 个模型（flux/qwen_image/mmdit/wan）的 forward 在结构 + 地标处调用 `TapRegistry` 的条件 `tap()`（`ctx->tap_registry`）——非请求即 no-op，不 pin + buffer、gallocr 无感。cache 需要读哪些张量由 per-substep 的 `SubstepPlan.taps` 决定， + 而非模型枚举字段。inject（复用注入）也走 registry：forward 在 `inject_at(i)` 处替换流 + 并跳到 `inject_resume()`。 +- **`CacheGraphScope` seam 已整块删除**（commit `cd0c9b7`，`-1667` 行）：旧的固定 `*_node` + 字段、`kCache*Name` 常量、`expand_cache_scope_nodes`、runner 的 `cache_scope_` 成员 + + `set_cache_scope`、`DiCacheGpuState` 独立 GPU 状态路径均已移除。只保留一个纯结果结构体 + `sd::DiffusionCacheResult`（host 读回的 output/feature/before/probe + DiCache 标量）。 +- **substep 是唯一路径**：`ED_CACHE_SUBSTEP` 门控消失，`run_branch` 无条件走 + `CacheGraphLowering::execute_substeps`（`cache_engine.cpp:196`）；旧 `execute` 及其 + hook 大 if/else 已删（commit `a2831c6`，`-444` 行）。8 个方法全部实现 `next_substep()`。 +- **DiCache probe 也已 tap 化**：checklist 曾标注「probe 未迁」的那项已完成——probe 的 + 跨步持久张量（prev_probe/probe_prev1/2）经 `TapRegistry::ProbeMetricOperands` 线程进 + indicator lowering，delta_y/delta_x/gamma 在图内归约后读回标量。 --- ## 3. 优点 -1. **解耦目标基本达成**：新增 cache 方法 = 写一个 `ICachePolicy`；新增模型 = 暴露一个 - `DiTModelCacheContract` + 一个 block-stack seam。摆脱了「模型数 × 方法数」的组合爆炸。 +1. **解耦目标已达成**：新增 cache 方法 = 写一个 `ICachePolicy`；新增模型 = 暴露一个 + `DiTModelCacheContract` + 在 forward 的结构地标处调用 `TapRegistry::tap()`（ModelIn/ + BlockOut[i]/ModelOut）。摆脱了「模型数 × 方法数」的组合爆炸，且模型侧不再 include 或 + 认识任何 cache 概念（见 §2.6）。 2. **失败是显式的**：SP-parallel 或无法切图时，Feature/Probe 方法经能力协商被正确拒绝， 而不是静默跑成 no-cache。 3. **CacheProgram 可 dump**（`ED_DUMP_CACHE_PROGRAM`），策略决策可观测。 @@ -93,137 +133,162 @@ vs 差距」。文中所有「未承重 / 死代码」类断言均已由代码 site。`cache_program.hpp` 的 `CachePhase::{FULL_ANCHOR,CORRECTION,REINTEGRATION}` 仅占位，实际只用到 `FORWARD/PROBE`（`redesign` 文档阶段三/四未做）。 - **SP-parallel 下 Feature/Probe 全失效**：`feature_cache_available()` 定义为 - `!can_attempt_graph_cut_segmented_compute()`（`ggml_extend.hpp:6309`），而后者在 - 「任何 process group 使能」时即为 true —— 故纯 SP 也会关掉 seam，SP 下只剩 Output - 粒度。这是引擎两大卖点（多卡并行 + 步缓存）当前无法叠加的直接体现。 -- **hook 耦合脆弱**：`gpu_metric`/`branch_key` 曾静默把所有 Feature cache（MagCache/ - TaylorSeer/SenCache）归零跑成 0/N（已修），说明 hook 表面 subtle、易回归。 + `!can_attempt_graph_cut_segmented_compute()`（`ggml_extend.hpp:6472-6473`），而后者在 + 「任何 process group 使能」时即为 true —— 故纯 SP 也会关掉 tap 捕获路径（pipeline 以 + `cache_seam_available = !cfg_parallel && feature_cache_available()` 传入 `CacheEngine::init`）， + SP 下只剩 Output 粒度。这是引擎两大卖点（多卡并行 + 步缓存）当前无法叠加的直接体现， + 也是解耦完成后**最大的剩余能力缺口**。 +- **~~hook 耦合脆弱~~（已随 dehook 消解）**：`gpu_metric`/`branch_key` 曾静默把所有 Feature + cache（MagCache/TaylorSeer/SenCache）归零跑成 0/N。该脆弱面源于旧 `CacheGraphScope` 的 + 固定 `*_node` 字段 + 双写路径；§2.6 的 dehook 已删除该 seam，tap 由模型按请求集主动登记， + 这一类回归的结构性成因已消除。保留此条作为历史教训：触及 tap/lowering 时仍按「skip 计数 + + PSNR≥地板」验收。 - **cache 路径曾有 GPU 内存泄漏**（已修），长跑仍需 chunked 生成绕过。 -### 4.2 脚手架层：搭好但未承重（均已 grep 核实） - -设计文档最有雄心的三块声明式机器，目前只落地了**决策层**： - -- **`CacheStateManager` 空转**：全库无 `read / write / rotate_history / read_history` - 调用（只有 `initialize / begin_step / commit_step` 被 engine 调用）；每个 policy 自持 - `std::unordered_map states_`（magcache/dicache/easycache/ucache/ - sencache/taylorseer/condition 均如此）。`commit_step / rollback_step` 是**显式 - no-op**（`cache_state_manager.cpp:105-116` 注释：宿主向量原地提交，无需 flush）—— - 所谓「事务式状态 + 回滚」并未提供任何真实安全性。典型未接驱动：TaylorSeer 声明 - `history_depth = n_derivatives + 1`（`taylorseer_policy.cpp:44`），正是 ring buffer - 的用武之地，却自己存 `TaylorSeerState` 的导数历史。 -- **`CacheOperatorRegistry` + 4 个 operator 是死代码**：`register_builtin_cache_operators` - **全库无人调用**，registry 从未实例化；`IdentityOperator/DifferenceOperator/ - LinearPredictOperator/WeightedBlendOperator` 的 `apply_host` 从未被执行。真正的残差/ - 预测/blend 数学内联在 policy（如 `magcache_policy.cpp:353` 直接读 `b.residual`）、 - lowering、pipeline 里。 -- **无 `GraphRepository` / 多静态图变体**：`GraphRepository`、`GraphVariantKey`、 +### 4.2 声明式机器：已承重 vs 仍是脚手架（2026-07-15 更新，均已 grep 核实） + +设计文档最有雄心的三块声明式机器，现状已从「只落地决策层」推进——**前两块已承重**： + +- **`CacheStateManager` 已承重（history ring + 设备槽后端）**：`CacheEngine::init` 把 + `state_` 传入 `CacheGraphLowering::execute_substeps`（`cache_engine.cpp:196`），lowering 的 + `ActionInterpreter` 真实调用 `state_.read_history` / `state_.write` / `state_.rotate_history` + （`cache_graph_lowering.cpp`），设备槽路径调 `state.read` + `state.alloc_device_entry`。 + TaylorSeer 的导数历史 ring、MagCache/Qwen 的**设备端单残差槽**都跑在其上。 + 仍是脚手架的部分：`commit_step`（`cache_state_manager.cpp`）/ `rollback_step` + 仍是**显式 no-op**——但这是**诚实的**：全库无 `rollback_step` 调用方，`commit_step` 只被 + `end_step` 无条件调用一次，且没有任何 cache 方法会中途 abort（失败走同步 full-compute + 回退，不涉及状态回滚）。所谓「事务式回滚安全」当前无触发点，是为未来细粒度/token-cache + 预留的接口占位，而非承重能力。各 policy 仍自持标量决策状态（`states_` map），这是**决策 + 逻辑**（该留在 policy），与 StateManager 管的**张量存储**已分离。 +- **`CacheOperatorRegistry` + operator 已承重**：`CacheEngine` 实例化 registry 并调 + `register_builtin_cache_operators`（`cache_engine.cpp:111`，此前死代码）；lowering 的 + DIFFERENCE/PREDICT/BLEND 动作经 `operators_.find(...)` 路由到 `DifferenceOperator/ + LinearPredictOperator/WeightedBlendOperator::apply_host`。 + EasyCache/UCache 的输出 diff、TaylorSeer 的历史外推 blend 都由 operator 执行。 +- **仍无 `GraphRepository` / 多静态图变体**：`GraphRepository`、`GraphVariantKey`、 `get_or_compile` 在 `src/`、`examples/` 及除 redesign 外的所有 docs 中**零命中**。 设计文档核心思想「运行时在多张预编译静态图变体间选择」**没有实现**。实际执行是 - hook/回调式：`cache_graph_lowering.cpp` 按 `variant->kind` 分支去调 - `hooks.full/capture/inject/inject_gpu/probe`。因此 `CacheProgram`/`GraphVariantPlan` - 是薄「决策描述符」（告诉 lowering 调哪个 hook），而非编译产物；`SegmentPlan`/ - `CacheAction` 字段大部分是未消费的元数据。 + tap 驱动的 runner pass 分派：`execute_substeps` 按 `SubstepPlan.op.kind` 分支去调 + `hooks.full` 与 tap 驱动的 `hooks.substep_capture/substep_probe/substep_inject_*`（device + 或 host）。因此 `CacheProgram`/`GraphVariantPlan` 是「决策描述符 + 动作序列」（告诉 lowering + 产出哪种 substep、执行哪些 slot/operator 动作），而非编译产物。 -> **小结**：声明式组件里，只有 `CacheProgram`/`GraphVariantPlan`（作为薄描述符）+ -> `CacheRunnerHooks` lowering 是承重的。`CacheStateManager` 的 history/事务 API 和 -> 整个 operator registry 是朝 redesign 目标搭建、但尚未承重的脚手架。 +> **小结**：声明式组件里，`CacheProgram`/`GraphVariantPlan`（描述符 + 动作）、 +> `CacheRunnerHooks` lowering、`CacheStateManager` 的 history/设备槽后端、`CacheOperatorRegistry` +> 均已承重——8 个方法（含 Flux/Qwen 的默认设备端 GPU 复用）真实跑在 StateManager+Operator 上； +> 且模型侧已彻底 dehook（§2.6）。仍是脚手架的只剩：`GraphRepository`/多变体编译（未建）， +> 以及 `commit_step`/`rollback_step` 的事务语义（诚实 no-op，无触发点，为未来预留）。 --- ## 5. 推荐前进目标：落地声明式引擎 **方向**：把 §4.2 的三块脚手架真正接上，兑现 `cache_framework_redesign.md` 的核心承诺。 +**现状（2026-07-14）**：前两层已落地承重，只剩第 3 层。 -### 5.1 具体含义（分层，每层可独立 byte-verify against 当前 hook 路径） +### 5.1 具体含义（分层，每层独立按「skip 计数 + PSNR≥地板」验收） -1. **StateManager 承重化**：把 `CacheStateManager` 接为 TaylorSeer / DiCache / MagCache - 的真实状态后端（history ring + 真正的 commit/rollback），替换各 policy 的私有 - `states_`。 -2. **Operator 化数学**：把残差 diff / 预测外推 / blend 路由到 `CacheOperator`，policy 只 - 引用 operator id，`register_builtin_cache_operators` 真正被调用。 -3. **GraphRepository + 变体编译**：引入真正的 FULL/REUSE/PREDICT 静态图变体编译与缓存， - 并与 `ED_CACHE_COMPILED_GRAPHS` 的 build-once 图复用融合（两者都想「少建图」）。 +1. **StateManager 承重化 ✅ 已完成**：`CacheStateManager` 已是 TaylorSeer / DiCache / + MagCache / Qwen 的真实状态后端（host history ring + 设备端单残差槽），替换了这些方法 + 张量存储层面的私有 `states_`。事务 commit/rollback 仍是诚实 no-op（无触发点，见 §4.2）。 +2. **Operator 化数学 ✅ 已完成**：残差 diff / 预测外推 / blend 已路由到 `CacheOperator`， + policy 引用 operator id，`register_builtin_cache_operators` 真正被调用。 +3. **GraphRepository + 变体编译 ⬜ 未做**：引入真正的 FULL/REUSE/PREDICT 静态图变体编译与 + 缓存，并与 `ED_CACHE_COMPILED_GRAPHS` 的 build-once 图复用融合（两者都想「少建图」）。 + 这是唯一剩下的声明式层，也是 SP 复合 / 细粒度的天然下游。 -### 5.2 优势 +### 5.2 已兑现的收益（层 1-2）与剩余价值（层 3） -- 兑现 redesign 承诺、偿还架构债； -- 为细粒度 / 轨迹修正 / token cache 打地基（这些方法都需要真实的多历史状态与 operator 组合）； -- 真正的事务状态让 rollback 安全（当前 no-op，异常步可能留下不一致历史）； +已兑现： +- 8 个方法 + Flux/Qwen 默认设备端 GPU 复用跑在 StateManager+Operator 上，`CacheProgram` + 不再只是决策描述符； - operator 化让新方法接入更快（组合已有 operator，而非每次内联新数学）； -- 与 compiled-graph 复用合流，可能进一步降低建图开销。 +- 设备槽把 skip 成本压到 ~1ms 且验证字节一致（见 §2.5、§5.5）。 + +剩余价值（层 3）： +- 为细粒度 / 轨迹修正 / token cache 打地基（需要真正的多变体静态图）； +- 与 compiled-graph 复用合流，可能进一步降低建图开销； +- 事务 rollback 安全当前**无触发点**，不是收益——除非未来出现会中途 abort 的方法。 ### 5.3 问题 / 风险（不回避） -- **当前 hook 路径已够用**：纯「接脚手架」若不解锁具体能力，可能是**无用户可见收益**的 - 工作量。因此本文主张——**把落地声明式引擎当作「手段」而非「目的」**：每一层都绑定一个 - 具体收益（如 rollback 安全性、或为 SP / 细粒度铺路），按可 byte-verify 的最小切片推进， - 拒绝为「架构纯洁」而重写。 -- **过度设计风险**：`GraphRepository` / 多变体在 GGML 静态图上与既有 `PlanCache` / +- **层 1-2 已证「脚手架可承重」**：设备槽验收字节一致（PSNR 100 dB）、skip 计数逐位一致， + 证明「接脚手架」确有具体收益（skip ~1ms、单一 GPU 复用路径）。这消解了原本「纯整洁化、 + 无用户可见收益」的担忧——本文仍主张**把落地当手段而非目的**，每层绑定具体收益。 +- **过度设计风险（层 3）**：`GraphRepository` / 多变体在 GGML 静态图上与既有 `PlanCache` / `compute_reuse`（`ED_CACHE_COMPILED_GRAPHS`）职责重叠。落地前须先厘清边界， 否则是两套「图缓存」互相打架（参见 `perf_gap_vs_diffusers.md` 的建图开销分析）。 -- **回归风险高**：触及最脆弱的 seam / graph-cut 子系统与曾出事的 hook 耦合（§4.1）。 -- **收益递延**：三层里只有第 3 层（GraphRepository）直接关联 SP 复合与细粒度这些 - 用户可感知能力；前两层是纯内部整洁化，收益偏「未来价值」。 +- **回归风险高**：层 3 触及最脆弱的 graph-cut / PlanCache 子系统。曾出事的 hook 耦合面 + 已随 §2.6 的 dehook 消解（seam 删除、tap 按请求集主动登记），但 tap/lowering 仍需谨慎， + 触及处按「skip 计数 + PSNR≥地板」验收。 +- **收益递延**：只有第 3 层（GraphRepository）直接关联 SP 复合与细粒度这些用户可感知能力。 -### 5.4 建议的首个最小切片 +### 5.4 首个最小切片（已完成，作为「脚手架可承重」的证明点） -只把 `CacheStateManager` 接为 **TaylorSeer** 的 history 后端——它有最明确的多历史驱动 -（`n_derivatives+1` 深度），改动面小、边界清晰。验收 = 输出对当前 hook 路径 **byte-identical**， -以此作为「脚手架可承重」的证明点，再决定是否推进第 2、3 层。**若这一层做完发现无实际收益， -即是「当前 hook 已够用」的有力信号，应及时止步而非硬推。** +原计划：只把 `CacheStateManager` 接为 **TaylorSeer** 的 history 后端（`n_derivatives+1` +深度），改动面小、边界清晰，作为证明点。**实际结果**：该切片 + 后续 7 个方法 + Flux/Qwen +默认设备端 GPU 复用全部落地并验证等价（skip 计数逐位一致 + PSNR≥地板；设备槽 vs host +更达字节一致 100 dB）。证明点成立——脚手架确可承重，且带来了具体收益（skip ~1ms、单一 +GPU 复用路径），故层 1-2 已推进完成。**下一步的判断点转移到层 3**：GraphRepository 与 +`ED_CACHE_COMPILED_GRAPHS` 职责重叠，落地前须先厘清边界（§5.3），否则宁可止步。 -### 5.5 实施进展（2026-07-13，声明式化已启动） +### 5.5 实施进展（2026-07-15，层 1-2 + 模型 dehook 已完成并验证） > **验收基准更正**：§5.1/§5.4 原写「byte-identical」，实测证伪——本引擎 CUDA 后端**逐帧 > 不确定**（同一 no-cache 配置连跑两次 PSNR ≈ 35 dB，非 100%）。因此等价性验收改为 > **①skip 计数逐位一致（精确复现，直接抓决策回归）+ ②输出 PSNR ≥ no-cache 自比的 > ~35 dB 噪声地板**。md5 比对是错误方法。 -**已落地并验证等价（低风险地基）：** - -- **能力接线**：`CacheEngine` 现实例化 `CacheOperatorRegistry` 并调用 - `register_builtin_cache_operators`（此前死代码），把 `state_` + `operators_` 传入 - `CacheGraphLowering::execute`。 -- **per-method 开关**：`variant_has_actions(program)` —— 发出 `CacheAction` 的方法走声明式 - `ActionInterpreter`（LOAD/STORE/DIFFERENCE/BLEND over slot + operator）；未迁移方法走 - 原 `reconstruct/observe` legacy 路径，逐字不变。可一次迁一个方法。 -- **EasyCache / UCache**（Output）：`make_output_diff_program`（FULL.after - DIFFERENCE→slot；REUSE.before LOAD+BLEND→输出）。skip 7/20、9/20 不变；diff 张量移入 slot。 -- **MagCache**（Feature，host 路径）：`make_feature_reuse_program`（FULL.after STORE 残差 - →slot；REUSE.before LOAD→注入）。host/GPU 均 reuse 12/20，host-decl vs gpu-legacy 33 dB。 - seam 的 capture/inject 保留为**合法的模型 compute 边界**，只把残差**存储**移入 slot。 -- **无回归**：DiCache 4/20、CacheDiT 1/20 不变；TaylorSeer 0/20 经 stash+重建证实为 - **既有**（默认 Flux 20 步 schedule 本就不触发），非本次迁移所致。 - -**剩余切片的风险评估（诚实）：** - -- **3b TaylorSeer 导数**：非机械移植。导数级联 `dy[d+1]=(dy[d]−dy_prev[d])/window` 是**有状态 - 递推**，且 window 随计算步间隔运行时变化；无法从 slot 里 K 个原始特征重建，也不匹配当前 - operator 的编译期 `params`。需要 IR 扩展（typed state-slot 或运行时参数绑定）。且默认设置 - skip 0/20，改动难以就地验证——建议连带 schedule 一起处理。 -- **3c GPU operator lowering**：需给 `ICacheOperator` 增加**发射 ggml 节点**的 lowering（而非 - host `apply_host`），并把 seam `cache_graph_scope.hpp` 的 `add_injected` 硬编码残差重建 - 改为 IR 驱动。触及最脆弱子系统，on-GPU 路径默认开启，高回归风险。 -- **3d DiCache probe/gamma**：`build_probe_metrics`/gamma 轨迹对齐迁到 IR，依赖 3c。 -- **Slice 4 彻底去 hook**：需为 flux/qwen/sd3/wan 各写 `IDiTGraphBuilder` 分段建图接口， - 让 seam 跨 graph-cut 存活——数周量级。 - -**结论**：声明式引擎已从「脚手架」转为「承重」——3 个方法真实跑在 StateManager+Operator 上， -`CacheProgram` 不再只是决策描述符。但「无 hook」的终态仍需上述 IR/seam 扩展；每一步继续按 -「skip 计数 + PSNR≥地板」验收，逐切片推进，触及 seam 处优先保留 fallback。 +**已落地并验证等价（8/8 方法 + 默认设备端 GPU 复用 + 模型 dehook）：** + +- **能力接线**：`CacheEngine` 实例化 `CacheOperatorRegistry` 并调用 + `register_builtin_cache_operators`（此前死代码），把 `state_` + `operators_` + 可选 + `ICacheDeviceStore` 传入 `CacheGraphLowering::execute_substeps`。 +- **统一 substep 路径**：8 个方法均实现 `next_substep()`，逐子步产出 `SubstepPlan`； + 中间层按 `SubstepOpKind` 分派到 tap 驱动的 runner pass。`ED_CACHE_SUBSTEP` 门控与旧 + `execute` 的 hook 大 if/else 已删（commit `a2831c6`）。 +- **Output（EasyCache/UCache/DBCache/CacheDiT）**：`OutputCompute`（FULL 后 DIFFERENCE→slot）/ + `OutputReuse`（LOAD+BLEND→输出）。skip 计数不变。 +- **Feature host 路径（MagCache/TaylorSeer/SenCache）**：`FeatureCompute`/`FeatureReuse` + （TaylorSeer 让 history ring 承重：raw-feature ring + ROTATE_HISTORY + 每步 reuse_coeffs + blend）。 +- **Probe（DiCache）**：`stop_after` 浅层前缀 + residual ring；probe 已 tap 化（delta_y/ + delta_x/gamma 经 `TapRegistry::ProbeMetricOperands` 在图内归约后读回标量）。 +- **默认设备端 GPU 复用（MagCache/Qwen，`ED_FEATURE_CACHE_GPU` 默认 ON）已承重并验证**： + `CacheStateManager` 拥有 per-slot 持久设备 buffer（`RunnerCacheDeviceStore`， + `ggml_extend.hpp:1866`），经 `substep_capture` d2d 存残差、`substep_inject_slot` 在图内 + `ggml_add(x_before, slot)` 复用。Flux **和 Qwen** 均已从 legacy `DiCacheGpuState` 路径 + 迁到设备槽。**验收（2026-07-14）**：设备槽 vs host-declarative **字节一致（PSNR 100 dB）**—— + Flux MagCache 30/50、Qwen MagCache 36/50，skip 计数逐位一致，两条路径 vs no-cache 分别 + 为 28.4 dB / 25.2 dB（与 benchmark 报告一致）。 +- **模型 dehook 已完成（§2.6）**：`CacheGraphScope` seam 整块删除（commit `cd0c9b7`）， + 4 个模型 forward 改为 `TapRegistry` 条件 tap，cache 层零依赖具体模型 header。 + +**剩余工作：** + +- **`GraphRepository` / 多静态图变体**：redesign 的第 3 层，未建。这是 SP + 缓存复合、 + 细粒度切点的天然下游。 +- **SP + 缓存复合**：SP 下 Feature/Probe 仍被能力协商拒绝（§4.1），当前只剩 Output 粒度。 + 解耦完成后这是**最大的用户可感知缺口**。 +- **事务语义（commit/rollback）**：诚实 no-op，无触发点（见 §4.2）。为未来预留，非当前债。 + +**结论**：声明式引擎已从「脚手架」转为「承重」——8 个方法 + Flux/Qwen 默认设备端 GPU 复用 +真实跑在 StateManager+Operator 上，`CacheProgram` 不再只是决策描述符；且**模型侧已彻底 dehook** +（原列为「数周量级终态」的 Slice 4 已完成，seam 删除，substep 是唯一路径）。仍是脚手架的只剩 +`GraphRepository`（未建）与事务语义（诚实 no-op）。每步继续按「skip 计数 + PSNR≥地板」验收 +（CUDA 逐帧不确定，md5 比对错误），触及 tap/lowering 处优先保留 fallback。 --- ## 6. 备选方向（简述，非推荐） - **SP + 缓存复合**：价值最高、最贴合引擎双卖点（多卡 + 步缓存当前无法叠加）。难点 = 让 - graph-cut planner 跨 segment 保留 seam 的命名中间输出——`ggml_extend.hpp:6306-6308` + graph-cut planner 跨 segment 保留 tap 的命名中间输出——`ggml_extend.hpp:6471` 明确记录了此冲突：「mid-graph capture is not preserved across segments」。这是「落地 声明式引擎」第 3 层（GraphRepository/变体编译）的天然下游，宜在其后做。Output 粒度缓存 - 理论上可较早搭 SP 便车（无需 seam）。 -- **巩固与加固**：删除或显式标注脚手架、修复脆弱 hook 耦合、刷新全量 benchmark - （当前 Qwen 行仍是旧数据）。低风险、不增能力，可作为任何方向的前置。 + 理论上可较早搭 SP 便车（无需 tap 捕获）。 +- **巩固与加固**：删除或显式标注剩余脚手架、刷新全量 benchmark（当前 Qwen 行仍是旧数据）。 + 低风险、不增能力，可作为任何方向的前置。（注：曾脆弱的 hook 耦合已随 §2.6 的 dehook 消解。） - **细粒度切点**：把拓扑从单一 `BLOCK_STACK` 扩到 per-block / attention·FFN / token 级， 解锁 partial-compute 与 token cache。研究前沿、上限高，但每个模型都要改 contract + seam， 测试矩阵组合爆炸。 From eeb3404856da48e431f6cc0d63cfe2945d981638 Mon Sep 17 00:00:00 2001 From: Liwei Lan Date: Wed, 15 Jul 2026 22:18:26 +0800 Subject: [PATCH 4/8] refactor: drive MagCache device capture through a cache-layer graph extension --- src/core/backend/ggml/ggml_extend.hpp | 36 +++++++++++++ .../cache/compile/cache_graph_lowering.cpp | 26 +++++++-- .../cache/compile/cache_graph_lowering.hpp | 12 +++-- .../optimization/cache/ir/graph_extension.hpp | 52 ++++++++++++++++++ .../optimization/cache/model/tap_registry.hpp | 10 ++++ src/dit_models/models/flux.hpp | 54 ++++++++++++------- src/dit_models/models/qwen_image.hpp | 44 +++++++++------ .../pipelines/flux_kontext_pipeline.cpp | 4 +- src/dit_models/pipelines/flux_pipeline.cpp | 4 +- .../pipelines/qwen_image_edit_pipeline.cpp | 4 +- .../pipelines/qwen_image_pipeline.cpp | 4 +- 11 files changed, 197 insertions(+), 53 deletions(-) create mode 100644 src/core/optimization/cache/ir/graph_extension.hpp diff --git a/src/core/backend/ggml/ggml_extend.hpp b/src/core/backend/ggml/ggml_extend.hpp index 43cd3958..c6bd5bc0 100644 --- a/src/core/backend/ggml/ggml_extend.hpp +++ b/src/core/backend/ggml/ggml_extend.hpp @@ -37,6 +37,7 @@ #include "optimization/cache/state/cache_device_store.hpp" #include "optimization/cache/model/tap_registry.hpp" #include "optimization/cache/compile/indicator_lowering.hpp" +#include "optimization/cache/operator/cache_operator.hpp" #include "edge-dit.h" #include "core/runtime/model_loader.h" @@ -4700,6 +4701,41 @@ struct GGMLRunner { expand_named(feat, "ed_cache_feature"); } } + // Cache-layer graph extensions: the cache lowering asked us to weave these + // operator nodes over the tapped tensors and pin the result. The runner + // executes them blindly — op->lower() emits the ggml nodes (e.g. a + // DIFFERENCE = ggml_sub), and this code does NOT know the math means + // "residual". Replaces the hardcoded capture_residual weave above for + // methods that have migrated (MagCache). The pass handoff then d2d-copies + // the pinned tensor into a device slot by output_name (CaptureToSlot). + for (const auto& ext : tap_registry_->extensions()) { + if (ext.op == nullptr) { + continue; + } + std::vector inputs; + bool inputs_ok = true; + for (const auto& a : ext.input_anchors) { + ggml_tensor* t = tap_registry_->get(a); + if (t == nullptr) { + inputs_ok = false; + break; + } + inputs.push_back(t); + } + if (!inputs_ok) { + continue; + } + for (ggml_tensor* t : ext.extra_inputs) { + inputs.push_back(t); + } + edgedit::cache::GraphLoweringContext gctx; + gctx.ctx = compute_ctx; + std::vector outputs; + if (ext.op->lower(gctx, inputs, ext.params, &outputs) && !outputs.empty() && + outputs[0] != nullptr) { + expand_named(outputs[0], ext.output_name); + } + } // Capture-writeback (DiCache device seed): also weave the probe residual // (BlockOut[m-1] - ModelIn) as "ed_cache_probe_resid" so the pass's // post-readback can d2d it into the runner's persistent probe-residual ring. diff --git a/src/core/optimization/cache/compile/cache_graph_lowering.cpp b/src/core/optimization/cache/compile/cache_graph_lowering.cpp index 84f51bd5..948d07b4 100644 --- a/src/core/optimization/cache/compile/cache_graph_lowering.cpp +++ b/src/core/optimization/cache/compile/cache_graph_lowering.cpp @@ -158,7 +158,6 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, const CacheRunnerHooks& hooks, CacheStateManager& state, const CacheOperatorRegistry& operators) { - (void)operators; policy.set_substep_input(hooks.input); policy.begin_substeps(step, condition_key); @@ -406,12 +405,29 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, // into the device slot when this plan writes it. const bool wants_capture = !plan.writes.empty(); if (device_slot && wants_capture && hooks.substep_capture) { - // Tap-driven capture: the runner requests - // ModelIn/ModelOut taps and weaves the residual — no CacheGraphScope. - auto alloc_slot = [&](const std::vector& residual_shape) -> void* { + // Tap-driven capture, cache-driven weave: the cache layer hands the + // runner a DIFFERENCE extension (model_out - model_in) and a slot to + // d2d the woven residual into. The runner requests the taps the + // extension references, weaves op->lower(), and pins the result — it + // never learns the math is a residual. Replaces the runner's old + // hardcoded ggml_sub + set_capture_residual seam. + GraphExtension cap; + cap.op = operators.find("cache.difference"); + cap.op_id = "cache.difference"; + // DifferenceOperator computes inputs[1] - inputs[0], so order is + // (model_in, model_out) => model_out - model_in, byte-identical to + // the old ggml_sub(mout, min). + cap.input_anchors = {AnchorRef::model_in(), AnchorRef::model_out()}; + cap.output_name = "ed_cache_feature"; + cap.sink = GraphExtension::Sink::CaptureToSlot; + cap.alloc_slot = [&](const std::vector& residual_shape) -> void* { return state.alloc_device_entry(condition_key, slot0, residual_shape); }; - y = hooks.substep_capture(alloc_slot); + std::vector exts; + if (cap.op != nullptr) { + exts.push_back(std::move(cap)); + } + y = hooks.substep_capture(std::move(exts)); if (!y.empty()) { CacheObservation obs; obs.kind = CacheObservation::Kind::Feature; diff --git a/src/core/optimization/cache/compile/cache_graph_lowering.hpp b/src/core/optimization/cache/compile/cache_graph_lowering.hpp index 3a25b151..a012558c 100644 --- a/src/core/optimization/cache/compile/cache_graph_lowering.hpp +++ b/src/core/optimization/cache/compile/cache_graph_lowering.hpp @@ -2,6 +2,7 @@ #include "core/optimization/cache/ir/cache_program.hpp" #include "core/optimization/cache/ir/runtime_decision.hpp" +#include "core/optimization/cache/ir/graph_extension.hpp" #include "core/optimization/cache/operator/cache_operator_registry.hpp" #include "core/optimization/cache/policy/cache_policy.hpp" #include "core/optimization/cache/state/cache_state_manager.hpp" @@ -24,11 +25,12 @@ struct CacheRunnerHooks { const sd::Tensor* input = nullptr; std::function()> full; - // Substep-path tap-driven capture (MagCache device slot): same contract as a - // device-slot store — alloc_slot returns the device slot for the residual - // shape; the runner drives it through the TapRegistry. - std::function( - const std::function&)>&)> substep_capture; + // Substep-path tap-driven capture (MagCache device slot): the lowering builds + // the GraphExtension list (a DIFFERENCE that weaves the residual and a slot to + // d2d it into); the runner requests the taps those extensions reference, runs + // the forward, weaves the operator nodes, and hands the pinned residual off to + // the slot. The runner never learns the math is a residual. + std::function(std::vector)> substep_capture; // Substep-path tap-driven DiCache device seed capture: full forward whose // post-readback refreshes the persistent DiCacheGpuState residual/probe rings diff --git a/src/core/optimization/cache/ir/graph_extension.hpp b/src/core/optimization/cache/ir/graph_extension.hpp new file mode 100644 index 00000000..74d86dd7 --- /dev/null +++ b/src/core/optimization/cache/ir/graph_extension.hpp @@ -0,0 +1,52 @@ +#pragma once + +#include +#include +#include + +#include "core/optimization/cache/ir/cache_action.hpp" // CacheOperatorParams +#include "core/optimization/cache/model/anchor.hpp" // AnchorRef + +// Forward declarations keep this IR header free of a real ggml.h include, so the +// declarative layer stays serializable and cheap to compile. Only the runner and +// the operator lowering (which dereference these) include the full definitions. +struct ggml_tensor; + +namespace edgedit { +namespace cache { + +class ICacheOperator; + +// A cache-layer instruction to EXTEND the model's compute graph: "run this +// operator over these tapped tensors and pin the result under output_name". The +// runner executes it blindly via op->lower() — it never learns what the math +// means (residual? blend?). This is what replaces the runner's hardcoded +// ggml_sub/ggml_add cache math: the cache layer now owns the graph nodes it wants +// woven, the model only reports structural landmarks (taps). +// +// Built by the lowering (which resolves `op` from the operator registry) and +// threaded to the runner through a substep hook. `op` points into the +// CacheEngine-owned registry, which outlives the pass. +struct GraphExtension { + // Where the woven output goes. + // CaptureToSlot : pinned under output_name, then d2d-copied into a device + // cache slot by the pass handoff (MagCache residual capture). + // ReplaceStream : the woven node replaces the block-stack stream at reuse + // (wired in the reuse step; not produced yet). + enum class Sink { CaptureToSlot, ReplaceStream }; + + const ICacheOperator* op = nullptr; // resolved operator (op->lower emits nodes) + std::string op_id; // operator id, for diagnostics + std::vector input_anchors; // structural taps feeding the operator, in order + std::vector extra_inputs; // non-tap device operands (slot / runtime scalar) + CacheOperatorParams params; // operator params (e.g. blend weights) + std::string output_name; // name to pin the woven node under (readback / d2d) + Sink sink = Sink::CaptureToSlot; + + // CaptureToSlot: allocate the device slot for the woven tensor's shape (known + // only post-capture). Mirrors the old alloc_slot hook contract. + std::function&)> alloc_slot; +}; + +} // namespace cache +} // namespace edgedit diff --git a/src/core/optimization/cache/model/tap_registry.hpp b/src/core/optimization/cache/model/tap_registry.hpp index 94137f51..dc710832 100644 --- a/src/core/optimization/cache/model/tap_registry.hpp +++ b/src/core/optimization/cache/model/tap_registry.hpp @@ -7,6 +7,7 @@ #include "core/optimization/cache/model/anchor.hpp" #include "core/optimization/cache/ir/indicator.hpp" +#include "core/optimization/cache/ir/graph_extension.hpp" struct ggml_tensor; @@ -66,6 +67,7 @@ class TapRegistry { recorded_.clear(); indicator_nodes_.clear(); indicators_.clear(); + extensions_.clear(); stop_after_block_ = -1; capture_residual_ = false; capture_writeback_ = false; @@ -94,6 +96,13 @@ class TapRegistry { void set_indicators(const std::vector& inds) { indicators_ = inds; } const std::vector& indicators() const { return indicators_; } + // ---- Cache-layer graph extensions. The cache lowering hands the runner a set + // of operator invocations to weave over the tapped tensors (residual capture, + // reuse blend). The runner executes them blindly via op->lower() — it does not + // know the math. Replaces the hardcoded ggml_sub residual weave. ---- + void set_extensions(std::vector exts) { extensions_ = std::move(exts); } + const std::vector& extensions() const { return extensions_; } + // Residual capture: when set, build_graph weaves (ModelOut - ModelIn) as a named // node "ed_cache_feature" so the pass can d2d-copy it into a device slot (the // block-stack residual for Feature/Probe reuse). Requires ModelIn + ModelOut @@ -191,6 +200,7 @@ class TapRegistry { std::unordered_map recorded_; std::vector indicator_nodes_; std::vector indicators_; + std::vector extensions_; int stop_after_block_ = -1; bool capture_residual_ = false; bool capture_writeback_ = false; diff --git a/src/dit_models/models/flux.hpp b/src/dit_models/models/flux.hpp index f7fe2918..90dfbddb 100644 --- a/src/dit_models/models/flux.hpp +++ b/src/dit_models/models/flux.hpp @@ -18,6 +18,7 @@ #include "dit_models/components/common/rope.hpp" #include "edge-dit.h" #include "core/runtime/model_loader.h" +#include "core/optimization/cache/ir/graph_extension.hpp" #include "parallel/sp_parallel.hpp" #define FLUX_GRAPH_SIZE 10240 @@ -4923,9 +4924,11 @@ namespace Flux { return std::move(pass.output); } - // ---- Substep-path tap-driven capture. Replaces the - // capture_to_slot hook: requests ModelIn/ModelOut taps, weaves the - // (ModelOut - ModelIn) residual, d2d-copies it into the device slot. ---- + // ---- Substep-path tap-driven capture. The cache layer hands us a set of + // GraphExtensions (a DIFFERENCE that weaves the residual + a slot to d2d it + // into); we request the taps they reference, run the forward, the base + // runner weaves the operator nodes, and we hand each CaptureToSlot result + // off to its device slot. The runner never learns the math is a residual. ---- sd::Tensor compute_substep_capture(int n_threads, const sd::Tensor& x, const sd::Tensor& timesteps, @@ -4935,24 +4938,37 @@ namespace Flux { const sd::Tensor& guidance, const std::vector>& ref_latents, bool increase_ref_index, - const std::function&)>& alloc_slot) { + std::vector extensions) { edgedit::cache::TapRegistry reg; - reg.set_requested({edgedit::cache::AnchorRef::model_in(), - edgedit::cache::AnchorRef::model_out()}); - reg.set_capture_residual(true); - std::function handoff = [&]() { - ggml_tensor* feat = get_cache_tensor_by_name("ed_cache_feature"); - if (feat == nullptr) { - return; + // Request exactly the taps the extensions reference (dedup preserved by + // set_requested's set semantics). + std::vector anchors; + for (const auto& ext : extensions) { + for (const auto& a : ext.input_anchors) { + anchors.push_back(a); } - std::vector shape; - const int nd = std::max(1, ggml_n_dims(feat)); - for (int i = 0; i < nd; ++i) { - shape.push_back(feat->ne[i]); - } - ggml_tensor* slot = static_cast(alloc_slot(shape)); - if (slot != nullptr && ggml_nbytes(slot) == ggml_nbytes(feat)) { - copy_named_cache_tensor_to("ed_cache_feature", slot); + } + reg.set_requested(anchors); + reg.set_extensions(extensions); + std::function handoff = [&]() { + for (const auto& ext : extensions) { + if (ext.sink != edgedit::cache::GraphExtension::Sink::CaptureToSlot || + !ext.alloc_slot) { + continue; + } + ggml_tensor* feat = get_cache_tensor_by_name(ext.output_name); + if (feat == nullptr) { + continue; + } + std::vector shape; + const int nd = std::max(1, ggml_n_dims(feat)); + for (int i = 0; i < nd; ++i) { + shape.push_back(feat->ne[i]); + } + ggml_tensor* slot = static_cast(ext.alloc_slot(shape)); + if (slot != nullptr && ggml_nbytes(slot) == ggml_nbytes(feat)) { + copy_named_cache_tensor_to(ext.output_name, slot); + } } }; auto get_graph = [&]() -> ggml_cgraph* { diff --git a/src/dit_models/models/qwen_image.hpp b/src/dit_models/models/qwen_image.hpp index 41b4dce2..b3b60695 100644 --- a/src/dit_models/models/qwen_image.hpp +++ b/src/dit_models/models/qwen_image.hpp @@ -20,6 +20,7 @@ #ifdef ED_ENABLE_CUDA_MODULATION #include "backend/ggml/ed_ggml_modulation_ext.hpp" #include "optimization/cache/compile/indicator_lowering.hpp" +#include "optimization/cache/ir/graph_extension.hpp" #endif namespace Qwen { @@ -3407,24 +3408,35 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* const sd::Tensor& context, const std::vector>& ref_latents, bool increase_ref_index, - const std::function&)>& alloc_slot) { + std::vector extensions) { edgedit::cache::TapRegistry reg; - reg.set_requested({edgedit::cache::AnchorRef::model_in(), - edgedit::cache::AnchorRef::model_out()}); - reg.set_capture_residual(true); - std::function handoff = [&]() { - ggml_tensor* feat = get_cache_tensor_by_name("ed_cache_feature"); - if (feat == nullptr) { - return; + std::vector anchors; + for (const auto& ext : extensions) { + for (const auto& a : ext.input_anchors) { + anchors.push_back(a); } - std::vector shape; - const int nd = std::max(1, ggml_n_dims(feat)); - for (int i = 0; i < nd; ++i) { - shape.push_back(feat->ne[i]); - } - ggml_tensor* slot = static_cast(alloc_slot(shape)); - if (slot != nullptr && ggml_nbytes(slot) == ggml_nbytes(feat)) { - copy_named_cache_tensor_to("ed_cache_feature", slot); + } + reg.set_requested(anchors); + reg.set_extensions(extensions); + std::function handoff = [&]() { + for (const auto& ext : extensions) { + if (ext.sink != edgedit::cache::GraphExtension::Sink::CaptureToSlot || + !ext.alloc_slot) { + continue; + } + ggml_tensor* feat = get_cache_tensor_by_name(ext.output_name); + if (feat == nullptr) { + continue; + } + std::vector shape; + const int nd = std::max(1, ggml_n_dims(feat)); + for (int i = 0; i < nd; ++i) { + shape.push_back(feat->ne[i]); + } + ggml_tensor* slot = static_cast(ext.alloc_slot(shape)); + if (slot != nullptr && ggml_nbytes(slot) == ggml_nbytes(feat)) { + copy_named_cache_tensor_to(ext.output_name, slot); + } } }; auto get_graph = [&]() -> ggml_cgraph* { diff --git a/src/dit_models/pipelines/flux_kontext_pipeline.cpp b/src/dit_models/pipelines/flux_kontext_pipeline.cpp index ab9fdcee..ae9cf489 100644 --- a/src/dit_models/pipelines/flux_kontext_pipeline.cpp +++ b/src/dit_models/pipelines/flux_kontext_pipeline.cpp @@ -1038,10 +1038,10 @@ bool FluxKontextPipeline::generate_one_image(const ed_image_generation_params_t* cache_runtime.granularity() == cache::CacheGranularity::Feature && Flux::FluxRunner::feature_gpu_enabled(); if (feature_gpu) { - hooks.substep_capture = [&](const std::function&)>& alloc_slot) { + hooks.substep_capture = [&](std::vector exts) { return flux_runner_->compute_substep_capture( n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, - cond_in.c_vector, guidance, ref_latents, false, alloc_slot); + cond_in.c_vector, guidance, ref_latents, false, std::move(exts)); }; hooks.substep_inject_slot = [&](void* slot, int region_start, int region_end) { return flux_runner_->compute_substep_inject_slot( diff --git a/src/dit_models/pipelines/flux_pipeline.cpp b/src/dit_models/pipelines/flux_pipeline.cpp index 884ce8e7..4920f8e9 100644 --- a/src/dit_models/pipelines/flux_pipeline.cpp +++ b/src/dit_models/pipelines/flux_pipeline.cpp @@ -949,10 +949,10 @@ bool FluxPipeline::generate_one_image(const ed_image_generation_params_t* params if (feature_gpu) { // Substep-path tap-driven capture: TapRegistry, // not CacheGraphScope. - hooks.substep_capture = [&](const std::function&)>& alloc_slot) { + hooks.substep_capture = [&](std::vector exts) { return flux_runner_->compute_substep_capture( n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, - cond_in.c_vector, guidance, {}, false, alloc_slot); + cond_in.c_vector, guidance, {}, false, std::move(exts)); }; // Substep-path tap-driven device inject (MagCache): registry inject. hooks.substep_inject_slot = [&](void* slot, int region_start, int region_end) { diff --git a/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp b/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp index 2a3161b2..75837b42 100644 --- a/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp +++ b/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp @@ -770,10 +770,10 @@ bool QwenImageEditPipeline::generate_one_image(const ed_image_generation_params_ cache_runtime.granularity() == cache::CacheGranularity::Feature && Qwen::QwenImageRunner::feature_gpu_enabled(); if (feature_gpu) { - hooks.substep_capture = [&, cond_in](const std::function&)>& alloc_slot) { + hooks.substep_capture = [&, cond_in](std::vector exts) { return diffusion_->compute_substep_capture( n_threads, x, timesteps, cond_in.c_crossattn, ref_latents, true, - alloc_slot); + std::move(exts)); }; hooks.substep_inject_slot = [&, cond_in](void* slot, int region_start, int region_end) { return diffusion_->compute_substep_inject_slot( diff --git a/src/dit_models/pipelines/qwen_image_pipeline.cpp b/src/dit_models/pipelines/qwen_image_pipeline.cpp index 8233797c..06d8b5a7 100644 --- a/src/dit_models/pipelines/qwen_image_pipeline.cpp +++ b/src/dit_models/pipelines/qwen_image_pipeline.cpp @@ -715,10 +715,10 @@ bool QwenImagePipeline::generate_one_image(const ed_image_generation_params_t* p if (feature_gpu) { // Substep-path tap-driven capture: same slot // contract, but driven through the TapRegistry, not CacheGraphScope. - hooks.substep_capture = [&, cond_in](const std::function&)>& alloc_slot) { + hooks.substep_capture = [&, cond_in](std::vector exts) { return diffusion_->compute_substep_capture( n_threads, x, timesteps, cond_in.c_crossattn, empty_ref_latents, false, - alloc_slot); + std::move(exts)); }; // Substep-path tap-driven device inject (MagCache): x_before + slot // via the forward's registry inject, no CacheGraphScope. From 372f959e2d207af963e080dda92ad156ec4b6ad7 Mon Sep 17 00:00:00 2001 From: Liwei Lan Date: Wed, 15 Jul 2026 23:06:19 +0800 Subject: [PATCH 5/8] refactor: drive MagCache device reuse through a cache-layer stream override --- src/core/backend/ggml/ggml_extend.hpp | 36 +++++++++++++++++++ .../cache/compile/cache_graph_lowering.cpp | 29 +++++++++++---- .../cache/compile/cache_graph_lowering.hpp | 11 +++--- .../optimization/cache/ir/graph_extension.hpp | 3 ++ .../optimization/cache/model/tap_registry.hpp | 22 ++++++++++++ src/dit_models/models/flux.hpp | 19 ++++++---- src/dit_models/models/qwen_image.hpp | 21 ++++++----- .../pipelines/flux_kontext_pipeline.cpp | 5 ++- src/dit_models/pipelines/flux_pipeline.cpp | 5 ++- .../pipelines/qwen_image_edit_pipeline.cpp | 4 +-- .../pipelines/qwen_image_pipeline.cpp | 4 +-- 11 files changed, 123 insertions(+), 36 deletions(-) diff --git a/src/core/backend/ggml/ggml_extend.hpp b/src/core/backend/ggml/ggml_extend.hpp index c6bd5bc0..c871b6cb 100644 --- a/src/core/backend/ggml/ggml_extend.hpp +++ b/src/core/backend/ggml/ggml_extend.hpp @@ -1857,6 +1857,42 @@ inline ggml_tensor* build_tap_inject(GGMLRunnerContext* ctx, ggml_tensor* x_befo } } +// Cache-layer stream override for substep reuse. The declarative counterpart to +// build_tap_inject: instead of the runner switching on a hardcoded InjectKind, the +// cache lowering hands over a ReplaceStream GraphExtension whose operator emits the +// reconstruction (x_before + residual, blend, etc.) via op->lower(). The runner +// stays blind to the math. x_before (the region-entry activation) is always the +// first operand; the ext's extra_inputs carry the cache operands (device slot, +// uploaded residual). Emits into whatever ggml_ctx the calling model's forward +// owns, so the same code weaves nodes into flux's / qwen's / any model's graph. +// Falls back to x_before (identity) when no ReplaceStream ext is present or the +// operator lowering fails, so an unserved reuse degrades to a no-op, never a crash. +inline ggml_tensor* build_stream_override(GGMLRunnerContext* ctx, ggml_tensor* x_before) { + if (ctx == nullptr || ctx->tap_registry == nullptr || x_before == nullptr) { + return x_before; + } + for (const auto& ext : ctx->tap_registry->extensions()) { + if (ext.sink != edgedit::cache::GraphExtension::Sink::ReplaceStream || + ext.op == nullptr) { + continue; + } + std::vector inputs; + inputs.push_back(x_before); + for (ggml_tensor* t : ext.extra_inputs) { + inputs.push_back(t); + } + edgedit::cache::GraphLoweringContext gctx; + gctx.ctx = ctx->ggml_ctx; + gctx.runtime_scalars = ext.runtime_scalars; + std::vector outputs; + if (ext.op->lower(gctx, inputs, ext.params, &outputs) && !outputs.empty() && + outputs[0] != nullptr) { + return outputs[0]; + } + } + return x_before; +} + // Device backing for CacheStateManager slots (implements the cache core's // abstract ICacheDeviceStore). Owns one persistent ggml_context + backend buffer, // allocated on the runtime backend OUTSIDE any per-step compute buffer, so slot diff --git a/src/core/optimization/cache/compile/cache_graph_lowering.cpp b/src/core/optimization/cache/compile/cache_graph_lowering.cpp index 948d07b4..b83f8e2e 100644 --- a/src/core/optimization/cache/compile/cache_graph_lowering.cpp +++ b/src/core/optimization/cache/compile/cache_graph_lowering.cpp @@ -376,14 +376,31 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, // whole-stack MagCache slice that is [0, end-of-stack) = [0, -1); a reuse // plan computes zero blocks but still injects over the whole stack. if (is_reuse && device_slot) { - // Serve the whole-stack residual from the persistent device slot: - // out = x_before + slot. The inject region is the whole stack [0, -1), - // NOT the (empty) compute range. Slot-not-ready falls through to a - // capture below. + // Serve the whole-stack residual from the persistent device slot as a + // cache-driven weave: out = WeightedBlend([x_before, slot], {1,1}) = + // x_before + slot. The lowering hands the runner a ReplaceStream + // extension; the model forward substitutes its output over the reuse + // region [0,-1) (whole stack) and never learns the math is a residual + // add. Slot-not-ready falls through to a capture below. CacheSlotHandle h = state.read(condition_key, slot0); if (h.valid && h.buffer != nullptr) { - // Tap-driven device inject (no CacheGraphScope). - y = hooks.substep_inject_slot(h.buffer, 0, -1); + GraphExtension inj; + inj.op = operators.find("cache.weighted_blend"); + inj.op_id = "cache.weighted_blend"; + // WeightedBlend: out = sum_k w[k]*inputs[k]. build_stream_override + // prepends x_before as inputs[0]; slot is inputs[1]. weights {1,1} + // => x_before + slot, byte-identical to the old build_tap_inject + // DeviceResidual branch (ggml_add(x_before, slot)). + inj.extra_inputs = {static_cast(h.buffer)}; + inj.params.floats = {1.0f, 1.0f}; + inj.sink = GraphExtension::Sink::ReplaceStream; + std::vector exts; + if (inj.op != nullptr) { + exts.push_back(std::move(inj)); + } + if (!exts.empty()) { + y = hooks.substep_inject_slot(std::move(exts)); + } } } else if (is_reuse && hooks.substep_inject_host) { // Host reuse (no device slot — Wan/SD3): the policy reconstructs the diff --git a/src/core/optimization/cache/compile/cache_graph_lowering.hpp b/src/core/optimization/cache/compile/cache_graph_lowering.hpp index a012558c..2ab5d276 100644 --- a/src/core/optimization/cache/compile/cache_graph_lowering.hpp +++ b/src/core/optimization/cache/compile/cache_graph_lowering.hpp @@ -52,10 +52,13 @@ struct CacheRunnerHooks { std::function substep_probe_host; // Host reuse (Wan): tap-driven inject of a host-reconstructed residual. (feature). std::function(const sd::Tensor&)> substep_inject_host; - // Device reuse tap-driven inject (Qwen/Flux), replacing inject_from_slot / - // inject_gpu on the substep path. slot: MagCache single residual; gpu: DiCache - // gamma-blend. Both drive the forward's registry inject (no CacheGraphScope). - std::function(void*, int, int)> substep_inject_slot; // (slot, start, end) + // Device reuse tap-driven inject (Qwen/Flux). MagCache: the lowering hands over + // a ReplaceStream GraphExtension (WeightedBlend of x_before + device slot); the + // runner requests the region taps, weaves op->lower(), and the forward + // substitutes it over the reuse region. The runner never learns the math is a + // residual add. substep_inject_gpu (DiCache gamma-blend) still uses the legacy + // hardcoded path pending its own migration. + std::function(std::vector)> substep_inject_slot; std::function(float, int, int)> substep_inject_gpu; // (gamma, start, end) }; diff --git a/src/core/optimization/cache/ir/graph_extension.hpp b/src/core/optimization/cache/ir/graph_extension.hpp index 74d86dd7..23988c24 100644 --- a/src/core/optimization/cache/ir/graph_extension.hpp +++ b/src/core/optimization/cache/ir/graph_extension.hpp @@ -39,6 +39,9 @@ struct GraphExtension { std::string op_id; // operator id, for diagnostics std::vector input_anchors; // structural taps feeding the operator, in order std::vector extra_inputs; // non-tap device operands (slot / runtime scalar) + std::vector runtime_scalars; // per-input broadcastable [1] device scalars + // (empty = use params.floats constants); DiCache + // gamma enters the blend here without a host copy CacheOperatorParams params; // operator params (e.g. blend weights) std::string output_name; // name to pin the woven node under (readback / d2d) Sink sink = Sink::CaptureToSlot; diff --git a/src/core/optimization/cache/model/tap_registry.hpp b/src/core/optimization/cache/model/tap_registry.hpp index dc710832..47eb8821 100644 --- a/src/core/optimization/cache/model/tap_registry.hpp +++ b/src/core/optimization/cache/model/tap_registry.hpp @@ -68,6 +68,9 @@ class TapRegistry { indicator_nodes_.clear(); indicators_.clear(); extensions_.clear(); + override_start_ = 0; + override_resume_ = -1; + override_active_ = false; stop_after_block_ = -1; capture_residual_ = false; capture_writeback_ = false; @@ -103,6 +106,22 @@ class TapRegistry { void set_extensions(std::vector exts) { extensions_ = std::move(exts); } const std::vector& extensions() const { return extensions_; } + // ---- ReplaceStream override region. Structural counterpart to the extensions: + // extensions() say WHAT stream replaces the region; this says WHICH region and + // WHERE to resume. The model forward skips the block interval [start, resume) + // and substitutes build_stream_override()'s output at the entry. The resume + // index is the model's own (it knows its block count); the cache only asks for + // a semantic whole-stack / sub-range. Kept separate from the legacy + // inject_* fields (DiCache DeviceBlend / Wan HostFeature still use those). ---- + void set_override_region(int start, int resume) { + override_start_ = start; + override_resume_ = resume; + override_active_ = true; + } + bool has_stream_override() const { return override_active_ && !extensions_.empty(); } + bool override_at(int i) const { return override_active_ && i == override_start_; } + int override_resume() const { return override_resume_; } + // Residual capture: when set, build_graph weaves (ModelOut - ModelIn) as a named // node "ed_cache_feature" so the pass can d2d-copy it into a device slot (the // block-stack residual for Feature/Probe reuse). Requires ModelIn + ModelOut @@ -201,6 +220,9 @@ class TapRegistry { std::vector indicator_nodes_; std::vector indicators_; std::vector extensions_; + int override_start_ = 0; + int override_resume_ = -1; + bool override_active_ = false; int stop_after_block_ = -1; bool capture_residual_ = false; bool capture_writeback_ = false; diff --git a/src/dit_models/models/flux.hpp b/src/dit_models/models/flux.hpp index 90dfbddb..7b52e28a 100644 --- a/src/dit_models/models/flux.hpp +++ b/src/dit_models/models/flux.hpp @@ -3582,7 +3582,7 @@ namespace Flux { // Tap-driven whole-stack inject (substep reuse): skip both block loops and // reconstruct x_before + residual via the registry. const bool tap_whole_inject = ctx->tap_registry != nullptr && - ctx->tap_registry->inject_active(); + ctx->tap_registry->has_stream_override(); const bool whole_inject = tap_whole_inject; ggml_tensor* cache_img_before = img; // Substep tap: block-stack input anchor (ModelIn). Conditional no-op @@ -4124,7 +4124,7 @@ namespace Flux { // (substep reuse): the loops ran zero blocks, so reconstruct // x_before + residual via the registry. if (tap_whole_inject) { - img = build_tap_inject(ctx, cache_img_before); + img = build_stream_override(ctx, cache_img_before); } // Substep tap: block-stack output anchor (ModelOut) — the residual's // "after" point (post-recombine, before final_layer). Conditional no-op @@ -4877,15 +4877,20 @@ namespace Flux { const sd::Tensor& guidance, const std::vector>& ref_latents, bool increase_ref_index, - ggml_tensor* slot, - int region_start, - int region_end) { - if (slot == nullptr) { + std::vector extensions) { + if (extensions.empty()) { return {}; } edgedit::cache::TapRegistry reg; + // Whole-stack reuse: the model owns the resume index (its own block + // count); the cache only asked for a semantic whole-stack override. The + // ReplaceStream extension carries WHAT stream replaces the region; the + // region [0, resume) carries WHERE. x_before is the model_in tap the + // forward passes to build_stream_override directly, so only the slot + // rides in the extension's extra_inputs. const int resume = flux_params.depth + flux_params.depth_single_blocks; - reg.set_inject_device_residual(slot, 0, resume); + reg.set_extensions(std::move(extensions)); + reg.set_override_region(0, resume); auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, {}); }; diff --git a/src/dit_models/models/qwen_image.hpp b/src/dit_models/models/qwen_image.hpp index b3b60695..d9347783 100644 --- a/src/dit_models/models/qwen_image.hpp +++ b/src/dit_models/models/qwen_image.hpp @@ -2978,9 +2978,9 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* // Tap-driven inject (substep reuse): at the region start, replace the // stream with the reconstructed x_before + residual and jump past the // region. - if (ctx->tap_registry != nullptr && ctx->tap_registry->inject_at(i)) { - img = build_tap_inject(ctx, cache_img_before); - i = ctx->tap_registry->inject_resume() - 1; + if (ctx->tap_registry != nullptr && ctx->tap_registry->override_at(i)) { + img = build_stream_override(ctx, cache_img_before); + i = ctx->tap_registry->override_resume() - 1; continue; } @@ -3352,15 +3352,18 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* const sd::Tensor& context, const std::vector>& ref_latents, bool increase_ref_index, - ggml_tensor* slot, - int region_start, - int region_end) { - if (slot == nullptr) { + std::vector extensions) { + if (extensions.empty()) { return {}; } edgedit::cache::TapRegistry reg; - const int resume = region_end < 0 ? qwen_image_params.num_layers : region_end; - reg.set_inject_device_residual(slot, region_start, resume); + // Region reuse: the model owns the resume index (num_layers = end of the + // single transformer-block stack). The cache asked for a whole-stack + // override; x_before is passed to build_stream_override directly by the + // forward, so only the device slot rides in the extension. + const int resume = qwen_image_params.num_layers; + reg.set_extensions(std::move(extensions)); + reg.set_override_region(0, resume); auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, ref_latents, increase_ref_index); }; diff --git a/src/dit_models/pipelines/flux_kontext_pipeline.cpp b/src/dit_models/pipelines/flux_kontext_pipeline.cpp index ae9cf489..0ecc4260 100644 --- a/src/dit_models/pipelines/flux_kontext_pipeline.cpp +++ b/src/dit_models/pipelines/flux_kontext_pipeline.cpp @@ -1043,11 +1043,10 @@ bool FluxKontextPipeline::generate_one_image(const ed_image_generation_params_t* n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, cond_in.c_vector, guidance, ref_latents, false, std::move(exts)); }; - hooks.substep_inject_slot = [&](void* slot, int region_start, int region_end) { + hooks.substep_inject_slot = [&](std::vector exts) { return flux_runner_->compute_substep_inject_slot( n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, - cond_in.c_vector, guidance, ref_latents, false, - static_cast(slot), region_start, region_end); + cond_in.c_vector, guidance, ref_latents, false, std::move(exts)); }; } if (cache_runtime.granularity() == cache::CacheGranularity::Probe) { diff --git a/src/dit_models/pipelines/flux_pipeline.cpp b/src/dit_models/pipelines/flux_pipeline.cpp index 4920f8e9..afc4afe4 100644 --- a/src/dit_models/pipelines/flux_pipeline.cpp +++ b/src/dit_models/pipelines/flux_pipeline.cpp @@ -955,11 +955,10 @@ bool FluxPipeline::generate_one_image(const ed_image_generation_params_t* params cond_in.c_vector, guidance, {}, false, std::move(exts)); }; // Substep-path tap-driven device inject (MagCache): registry inject. - hooks.substep_inject_slot = [&](void* slot, int region_start, int region_end) { + hooks.substep_inject_slot = [&](std::vector exts) { return flux_runner_->compute_substep_inject_slot( n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, - cond_in.c_vector, guidance, {}, false, - static_cast(slot), region_start, region_end); + cond_in.c_vector, guidance, {}, false, std::move(exts)); }; } if (cache_runtime.granularity() == cache::CacheGranularity::Probe) { diff --git a/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp b/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp index 75837b42..a21d60ad 100644 --- a/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp +++ b/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp @@ -775,10 +775,10 @@ bool QwenImageEditPipeline::generate_one_image(const ed_image_generation_params_ n_threads, x, timesteps, cond_in.c_crossattn, ref_latents, true, std::move(exts)); }; - hooks.substep_inject_slot = [&, cond_in](void* slot, int region_start, int region_end) { + hooks.substep_inject_slot = [&, cond_in](std::vector exts) { return diffusion_->compute_substep_inject_slot( n_threads, x, timesteps, cond_in.c_crossattn, ref_latents, true, - static_cast(slot), region_start, region_end); + std::move(exts)); }; } if (cache_runtime.granularity() == cache::CacheGranularity::Probe) { diff --git a/src/dit_models/pipelines/qwen_image_pipeline.cpp b/src/dit_models/pipelines/qwen_image_pipeline.cpp index 06d8b5a7..c17deb6a 100644 --- a/src/dit_models/pipelines/qwen_image_pipeline.cpp +++ b/src/dit_models/pipelines/qwen_image_pipeline.cpp @@ -722,10 +722,10 @@ bool QwenImagePipeline::generate_one_image(const ed_image_generation_params_t* p }; // Substep-path tap-driven device inject (MagCache): x_before + slot // via the forward's registry inject, no CacheGraphScope. - hooks.substep_inject_slot = [&, cond_in](void* slot, int region_start, int region_end) { + hooks.substep_inject_slot = [&, cond_in](std::vector exts) { return diffusion_->compute_substep_inject_slot( n_threads, x, timesteps, cond_in.c_crossattn, empty_ref_latents, false, - static_cast(slot), region_start, region_end); + std::move(exts)); }; } if (cache_runtime.granularity() == cache::CacheGranularity::Probe) { From 4506341400ffeb28d71921a2217ff1ad0ec2e087 Mon Sep 17 00:00:00 2001 From: Liwei Lan Date: Wed, 15 Jul 2026 23:37:47 +0800 Subject: [PATCH 6/8] refactor: drive DiCache device reuse through a cache-layer gamma-blend operator --- src/core/backend/ggml/ggml_extend.hpp | 9 +--- .../cache/compile/cache_graph_lowering.cpp | 21 ++++++-- .../cache/compile/cache_graph_lowering.hpp | 6 ++- .../cache/operator/cache_operators.cpp | 53 +++++++++++++++++++ src/dit_models/models/flux.hpp | 20 +++---- src/dit_models/models/qwen_image.hpp | 22 ++++---- .../pipelines/flux_kontext_pipeline.cpp | 5 +- src/dit_models/pipelines/flux_pipeline.cpp | 5 +- .../pipelines/qwen_image_edit_pipeline.cpp | 5 +- .../pipelines/qwen_image_pipeline.cpp | 5 +- 10 files changed, 104 insertions(+), 47 deletions(-) diff --git a/src/core/backend/ggml/ggml_extend.hpp b/src/core/backend/ggml/ggml_extend.hpp index c871b6cb..05b8e9bd 100644 --- a/src/core/backend/ggml/ggml_extend.hpp +++ b/src/core/backend/ggml/ggml_extend.hpp @@ -1845,13 +1845,8 @@ inline ggml_tensor* build_tap_inject(GGMLRunnerContext* ctx, ggml_tensor* x_befo return ggml_add(c, x_before, reg->inject_input()); case TapRegistry::InjectKind::DeviceResidual: return ggml_add(c, x_before, reg->inject_resid1()); - case TapRegistry::InjectKind::DeviceBlend: { - // x_before + resid2 + gamma*(resid1 - resid2) (matches add_injected). - ggml_tensor* diff = ggml_sub(c, reg->inject_resid1(), reg->inject_resid2()); - ggml_tensor* scaled = ggml_mul(c, diff, reg->inject_gamma()); - ggml_tensor* aligned = ggml_add(c, reg->inject_resid2(), scaled); - return ggml_add(c, x_before, aligned); - } + // DeviceBlend (DiCache gamma-blend) migrated to the cache.gamma_blend + // operator via build_stream_override; no longer served here. default: return x_before; } diff --git a/src/core/optimization/cache/compile/cache_graph_lowering.cpp b/src/core/optimization/cache/compile/cache_graph_lowering.cpp index b83f8e2e..803bb477 100644 --- a/src/core/optimization/cache/compile/cache_graph_lowering.cpp +++ b/src/core/optimization/cache/compile/cache_graph_lowering.cpp @@ -310,11 +310,22 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, // ---- Extrapolate reuse (DiCache): reconstruct on-device from the residual // ring using the clamped gamma the probe produced. Zero compute. ---- if (plan.op.kind == SubstepOpKind::Extrapolate) { - if (!plan.op.coeffs.empty()) { - const float gamma = plan.op.coeffs.front(); - // Tap-driven device inject (no CacheGraphScope). - if (hooks.substep_inject_gpu) { - y = hooks.substep_inject_gpu(gamma, 0, -1); + if (!plan.op.coeffs.empty() && hooks.substep_inject_gpu) { + // Cache-driven gamma-blend: gamma is the host constant the probe + // produced (already clamped by the policy), baked into params.floats. + // The model fills resid1/resid2 from its ring; build_stream_override + // weaves x_before + resid2 + gamma*(resid1-resid2) via op->lower(). + GraphExtension inj; + inj.op = operators.find("cache.gamma_blend"); + inj.op_id = "cache.gamma_blend"; + inj.params.floats = {plan.op.coeffs.front()}; + inj.sink = GraphExtension::Sink::ReplaceStream; + std::vector exts; + if (inj.op != nullptr) { + exts.push_back(std::move(inj)); + } + if (!exts.empty()) { + y = hooks.substep_inject_gpu(std::move(exts)); } } if (y.empty()) { diff --git a/src/core/optimization/cache/compile/cache_graph_lowering.hpp b/src/core/optimization/cache/compile/cache_graph_lowering.hpp index 2ab5d276..76696aad 100644 --- a/src/core/optimization/cache/compile/cache_graph_lowering.hpp +++ b/src/core/optimization/cache/compile/cache_graph_lowering.hpp @@ -59,7 +59,11 @@ struct CacheRunnerHooks { // residual add. substep_inject_gpu (DiCache gamma-blend) still uses the legacy // hardcoded path pending its own migration. std::function(std::vector)> substep_inject_slot; - std::function(float, int, int)> substep_inject_gpu; // (gamma, start, end) + // DiCache device reuse: the lowering hands over a ReplaceStream GraphExtension + // (cache.gamma_blend, gamma baked into params.floats as a host constant); the + // model fills the residual-ring operands and build_stream_override weaves the + // blend. The runner never learns the math is a gamma extrapolation. + std::function(std::vector)> substep_inject_gpu; }; // Drives the policy's substep loop against the runner hooks, using the policy for diff --git a/src/core/optimization/cache/operator/cache_operators.cpp b/src/core/optimization/cache/operator/cache_operators.cpp index 6ea42e92..96e1b09e 100644 --- a/src/core/optimization/cache/operator/cache_operators.cpp +++ b/src/core/optimization/cache/operator/cache_operators.cpp @@ -211,6 +211,58 @@ class WeightedBlendOperator final : public ICacheOperator { } }; +// out = x_before + resid2 + gamma*(resid1 - resid2). inputs: [x_before, resid1, +// resid2]; gamma is a host-known compile-time constant in params.floats[0] (the +// policy already clamped it), emitted via ggml_scale — DiCache reuse never reads +// gamma back to host, and since the value is known at graph-build time it need not +// be an on-device tensor. Matches the legacy build_tap_inject DeviceBlend weave +// (sub -> scale -> add -> add), modulo scale-vs-mul kernel choice. +class GammaBlendOperator final : public ICacheOperator { +public: + CacheOperatorSchema schema() const override { return {"cache.gamma_blend", 3, 3, 1, true}; } + + bool apply_host(const std::vector*>& inputs, + const CacheOperatorParams& params, + std::vector>* outputs) const override { + if (inputs.size() != 3 || inputs[0] == nullptr || inputs[1] == nullptr || + inputs[2] == nullptr || params.floats.empty() || outputs == nullptr) { + return false; + } + const int64_t n = inputs[0]->numel(); + if (inputs[1]->numel() != n || inputs[2]->numel() != n) { + return false; + } + const float g = params.floats[0]; + sd::Tensor out(inputs[0]->shape()); + const float* xb = inputs[0]->data(); + const float* r1 = inputs[1]->data(); + const float* r2 = inputs[2]->data(); + float* o = out.data(); + for (int64_t i = 0; i < n; ++i) { + o[i] = xb[i] + r2[i] + g * (r1[i] - r2[i]); + } + outputs->assign(1, std::move(out)); + return true; + } + + // sub(resid1,resid2) -> scale(diff,gamma) -> add(resid2,·) -> add(x_before,·). + bool lower(GraphLoweringContext& ctx, + const std::vector& inputs, + const CacheOperatorParams& params, + std::vector* outputs) const override { + if (inputs.size() != 3 || inputs[0] == nullptr || inputs[1] == nullptr || + inputs[2] == nullptr || ctx.ctx == nullptr || params.floats.empty() || + outputs == nullptr) { + return false; + } + ggml_tensor* diff = ggml_sub(ctx.ctx, inputs[1], inputs[2]); // resid1 - resid2 + ggml_tensor* scaled = ggml_scale(ctx.ctx, diff, params.floats[0]); // gamma*(...) + ggml_tensor* aligned = ggml_add(ctx.ctx, inputs[2], scaled); // resid2 + ... + outputs->assign(1, ggml_add(ctx.ctx, inputs[0], aligned)); // x_before + ... + return true; + } +}; + } // namespace void register_builtin_cache_operators(CacheOperatorRegistry* registry) { @@ -223,6 +275,7 @@ void register_builtin_cache_operators(CacheOperatorRegistry* registry) { registry->register_operator(std::make_unique()); registry->register_operator(std::make_unique()); registry->register_operator(std::make_unique()); + registry->register_operator(std::make_unique()); // cache.history_rotate is a state-manager action, not a tensor op; it has no // host lowering and is handled directly by CacheStateManager::rotate_history. } diff --git a/src/dit_models/models/flux.hpp b/src/dit_models/models/flux.hpp index 7b52e28a..41410600 100644 --- a/src/dit_models/models/flux.hpp +++ b/src/dit_models/models/flux.hpp @@ -4359,7 +4359,6 @@ namespace Flux { SDVersion version; bool use_mask = false; sd::Tensor inject_feature_host_; // kept alive across cache inject build - sd::Tensor gamma_scalar_host_; // GPU DiCache: [1] gamma for reconstruct // ---- Persistent cross-step GPU state for DiCache (Probe granularity) ---- // Holds the last computed step's probe state, block input, probe residuals @@ -4907,22 +4906,23 @@ namespace Flux { const sd::Tensor& guidance, const std::vector>& ref_latents, bool increase_ref_index, - float gamma, - const void* branch_key, - int region_start, - int region_end) { + std::vector extensions, + const void* branch_key) { auto it = dicache_gpu_states_.find(branch_key); if (it == dicache_gpu_states_.end() || it->second.buffer == nullptr || - it->second.resid_count < 2) { - return {}; // not enough history yet + it->second.resid_count < 2 || extensions.empty()) { + return {}; // not enough history yet -> lowering falls back to full } DiCacheGpuState& s = it->second; edgedit::cache::TapRegistry reg; - gamma_scalar_host_ = sd::Tensor({1}, std::vector{gamma}); const int resume = flux_params.depth + flux_params.depth_single_blocks; + // Shallow migration: gamma is baked into the extension by the lowering; + // the model only supplies the residual-ring device operands. build_graph + // prepends x_before, so gamma_blend sees [x_before, resid_prev1, resid_prev2]. + extensions[0].extra_inputs = {s.resid_prev1, s.resid_prev2}; + reg.set_extensions(std::move(extensions)); + reg.set_override_region(0, resume); auto get_graph = [&]() -> ggml_cgraph* { - ggml_tensor* g = make_input(gamma_scalar_host_); - reg.set_inject_device_blend(s.resid_prev1, s.resid_prev2, g, 0, resume); return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, {}); }; auto pass = run_substep_pass(get_graph, n_threads, ®, x.dim(), {}); diff --git a/src/dit_models/models/qwen_image.hpp b/src/dit_models/models/qwen_image.hpp index d9347783..3c8017e9 100644 --- a/src/dit_models/models/qwen_image.hpp +++ b/src/dit_models/models/qwen_image.hpp @@ -3095,7 +3095,6 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* std::vector modulate_index_vec; SDVersion version; sd::Tensor inject_feature_host_; // kept alive across cache inject build - sd::Tensor gamma_scalar_host_; // GPU DiCache: [1] gamma for reconstruct // ---- Persistent cross-step GPU state for DiCache (Probe granularity) ---- // Mirrors FluxRunner::DiCacheGpuState: holds the last computed step's probe @@ -3372,29 +3371,28 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* } // DiCache device reuse: x_before + resid2 + gamma*(resid1-resid2) from the - // persistent 2-deep ring, gamma uploaded as a [1] input. + // persistent 2-deep ring. Shallow migration: gamma is baked into the + // cache extension (host constant); the model only fills the ring operands. sd::Tensor compute_substep_inject_gpu(int n_threads, const sd::Tensor& x, const sd::Tensor& timesteps, const sd::Tensor& context, const std::vector>& ref_latents, bool increase_ref_index, - float gamma, - const void* branch_key, - int region_start, - int region_end) { + std::vector extensions, + const void* branch_key) { auto it = dicache_gpu_states_.find(branch_key); if (it == dicache_gpu_states_.end() || it->second.buffer == nullptr || - it->second.resid_count < 2) { - return {}; // not enough history yet + it->second.resid_count < 2 || extensions.empty()) { + return {}; // not enough history yet -> lowering falls back to full } DiCacheGpuState& s = it->second; edgedit::cache::TapRegistry reg; - gamma_scalar_host_ = sd::Tensor({1}, std::vector{gamma}); - const int resume = region_end < 0 ? qwen_image_params.num_layers : region_end; + const int resume = qwen_image_params.num_layers; + extensions[0].extra_inputs = {s.resid_prev1, s.resid_prev2}; + reg.set_extensions(std::move(extensions)); + reg.set_override_region(0, resume); auto get_graph = [&]() -> ggml_cgraph* { - ggml_tensor* g = make_input(gamma_scalar_host_); - reg.set_inject_device_blend(s.resid_prev1, s.resid_prev2, g, region_start, resume); return build_graph(x, timesteps, context, ref_latents, increase_ref_index); }; auto pass = run_substep_pass(get_graph, n_threads, ®, x.dim(), {}); diff --git a/src/dit_models/pipelines/flux_kontext_pipeline.cpp b/src/dit_models/pipelines/flux_kontext_pipeline.cpp index 0ecc4260..5be6ef7b 100644 --- a/src/dit_models/pipelines/flux_kontext_pipeline.cpp +++ b/src/dit_models/pipelines/flux_kontext_pipeline.cpp @@ -1058,11 +1058,10 @@ bool FluxKontextPipeline::generate_one_image(const ed_image_generation_params_t* delta_minus); }; if (Flux::FluxRunner::dicache_gpu_enabled()) { - hooks.substep_inject_gpu = [&, branch_key](float gamma, int region_start, int region_end) { + hooks.substep_inject_gpu = [&, branch_key](std::vector exts) { return flux_runner_->compute_substep_inject_gpu(n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, cond_in.c_vector, - guidance, ref_latents, false, gamma, branch_key, - region_start, region_end); + guidance, ref_latents, false, std::move(exts), branch_key); }; const int probe_depth = cache_runtime.dicache_probe_depth(); hooks.substep_capture_probe = [&, branch_key, probe_depth]() { diff --git a/src/dit_models/pipelines/flux_pipeline.cpp b/src/dit_models/pipelines/flux_pipeline.cpp index afc4afe4..4ec14ed0 100644 --- a/src/dit_models/pipelines/flux_pipeline.cpp +++ b/src/dit_models/pipelines/flux_pipeline.cpp @@ -977,11 +977,10 @@ bool FluxPipeline::generate_one_image(const ed_image_generation_params_t* params // ring blend), instead of the legacy on-device reconstruction. if (Flux::FluxRunner::dicache_gpu_enabled()) { // Substep-path tap-driven device inject (DiCache gamma-blend). - hooks.substep_inject_gpu = [&, branch_key](float gamma, int region_start, int region_end) { + hooks.substep_inject_gpu = [&, branch_key](std::vector exts) { return flux_runner_->compute_substep_inject_gpu(n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, cond_in.c_vector, - guidance, {}, false, gamma, branch_key, - region_start, region_end); + guidance, {}, false, std::move(exts), branch_key); }; // Substep-path tap-driven seed capture: full forward that // refreshes the DiCacheGpuState rings device-to-device (replaces diff --git a/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp b/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp index a21d60ad..62081e32 100644 --- a/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp +++ b/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp @@ -789,10 +789,9 @@ bool QwenImageEditPipeline::generate_one_image(const ed_image_generation_params_ true, depth, branch_key, delta_minus); }; if (Qwen::QwenImageRunner::dicache_gpu_enabled()) { - hooks.substep_inject_gpu = [&, cond_in, branch_key](float gamma, int region_start, int region_end) { + hooks.substep_inject_gpu = [&, cond_in, branch_key](std::vector exts) { return diffusion_->compute_substep_inject_gpu(n_threads, x, timesteps, cond_in.c_crossattn, - ref_latents, true, gamma, branch_key, - region_start, region_end); + ref_latents, true, std::move(exts), branch_key); }; const int probe_depth = cache_runtime.dicache_probe_depth(); hooks.substep_capture_probe = [&, cond_in, branch_key, probe_depth]() { diff --git a/src/dit_models/pipelines/qwen_image_pipeline.cpp b/src/dit_models/pipelines/qwen_image_pipeline.cpp index c17deb6a..39ad9019 100644 --- a/src/dit_models/pipelines/qwen_image_pipeline.cpp +++ b/src/dit_models/pipelines/qwen_image_pipeline.cpp @@ -742,10 +742,9 @@ bool QwenImagePipeline::generate_one_image(const ed_image_generation_params_t* p // host probe path instead (residual-ring blend). if (Qwen::QwenImageRunner::dicache_gpu_enabled()) { // Substep-path tap-driven device inject (DiCache gamma-blend). - hooks.substep_inject_gpu = [&, cond_in, branch_key](float gamma, int region_start, int region_end) { + hooks.substep_inject_gpu = [&, cond_in, branch_key](std::vector exts) { return diffusion_->compute_substep_inject_gpu(n_threads, x, timesteps, cond_in.c_crossattn, - empty_ref_latents, false, gamma, branch_key, - region_start, region_end); + empty_ref_latents, false, std::move(exts), branch_key); }; // Substep-path tap-driven seed capture: full forward that // refreshes the DiCacheGpuState rings device-to-device (replaces From 7664b6f246f3153e755d4f43c566c12f282068a0 Mon Sep 17 00:00:00 2001 From: Liwei Lan Date: Thu, 16 Jul 2026 00:37:38 +0800 Subject: [PATCH 7/8] refactor: drive DiCache probe metrics through cache-layer indicator operators --- src/core/backend/ggml/ggml_extend.hpp | 34 +----- .../cache/compile/cache_graph_lowering.cpp | 7 +- .../cache/compile/cache_graph_lowering.hpp | 10 +- .../optimization/cache/ir/graph_extension.hpp | 6 +- .../cache/operator/cache_operators.cpp | 104 ++++++++++++++++++ src/dit_models/models/flux.hpp | 49 +++++++-- src/dit_models/models/qwen_image.hpp | 44 ++++++-- .../pipelines/flux_kontext_pipeline.cpp | 4 +- src/dit_models/pipelines/flux_pipeline.cpp | 4 +- .../pipelines/qwen_image_edit_pipeline.cpp | 4 +- .../pipelines/qwen_image_pipeline.cpp | 4 +- 11 files changed, 206 insertions(+), 64 deletions(-) diff --git a/src/core/backend/ggml/ggml_extend.hpp b/src/core/backend/ggml/ggml_extend.hpp index 05b8e9bd..8341471e 100644 --- a/src/core/backend/ggml/ggml_extend.hpp +++ b/src/core/backend/ggml/ggml_extend.hpp @@ -4781,36 +4781,10 @@ struct GGMLRunner { expand_named(ggml_sub(compute_ctx, probe, before), "ed_cache_probe_resid"); } } - // DiCache probe metrics (delta_y/delta_x/gamma): weave the exact - // build_probe_metrics form from the tapped probe (BlockOut[m]) + before - // (ModelIn) anchors and the runner-owned persistent operands. Only the - // scalars are read back. - if (tap_registry_->probe_metrics()) { - ggml_tensor* probe = tap_registry_->get(tap_registry_->probe_anchor()); - ggml_tensor* before = tap_registry_->get(tap_registry_->before_anchor()); - const auto& ops = tap_registry_->probe_ops(); - auto rel = [&](ggml_tensor* cur, ggml_tensor* ref) -> ggml_tensor* { - ggml_tensor* num = ggml_sum(compute_ctx, ggml_abs(compute_ctx, ggml_sub(compute_ctx, cur, ref))); - ggml_tensor* den = ggml_sum(compute_ctx, ggml_abs(compute_ctx, ref)); - return ggml_div(compute_ctx, num, den); - }; - if (probe != nullptr && before != nullptr) { - if (ops.prev_probe != nullptr) { - expand_named(rel(probe, ops.prev_probe), "cache_ind:delta_y"); - } - if (ops.want_delta_x && ops.prev_input != nullptr) { - expand_named(rel(before, ops.prev_input), "cache_ind:delta_x"); - } - if (ops.want_gamma && ops.probe_prev1 != nullptr && ops.probe_prev2 != nullptr) { - ggml_tensor* cur_resid = ggml_sub(compute_ctx, probe, before); - ggml_tensor* num = ggml_sum(compute_ctx, ggml_abs(compute_ctx, - ggml_sub(compute_ctx, cur_resid, ops.probe_prev2))); - ggml_tensor* den = ggml_sum(compute_ctx, ggml_abs(compute_ctx, - ggml_sub(compute_ctx, ops.probe_prev1, ops.probe_prev2))); - expand_named(ggml_div(compute_ctx, num, den), "cache_ind:gamma"); - } - } - } + // DiCache probe metrics (delta_y/delta_x/gamma) are now woven by cache + // operators (cache.rel_l1 / cache.gamma_indicator) via the extensions loop + // above — the model builds Indicator-sink extensions in compute_substep_probe. + // The hardcoded reduction weave that lived here has been retired. // Also pin the raw taps (so a method that reads an anchor tensor directly // can). Indicator/feature operands are already expanded above. for (const auto& kv : tap_registry_->recorded()) { diff --git a/src/core/optimization/cache/compile/cache_graph_lowering.cpp b/src/core/optimization/cache/compile/cache_graph_lowering.cpp index 803bb477..56db5d19 100644 --- a/src/core/optimization/cache/compile/cache_graph_lowering.cpp +++ b/src/core/optimization/cache/compile/cache_graph_lowering.cpp @@ -290,9 +290,10 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, ? plan.blocks.end : std::max(1, plan.blocks.begin); if (hooks.substep_probe) { - // Tap-driven probe: delta_y/gamma computed - // on-device from taps + persistent operands, no CacheGraphScope. - sd::DiffusionCacheResult pr = hooks.substep_probe(depth); + // Tap-driven probe: delta_y/delta_x/gamma woven by cache operators + // (rel_l1 / gamma_indicator), resolved from the registry; the model + // supplies the probe-history device operands. + sd::DiffusionCacheResult pr = hooks.substep_probe(depth, operators); result.indicators["delta_y"] = pr.delta_y; result.indicators["delta_x"] = pr.delta_x; result.indicators["gamma"] = pr.gamma; diff --git a/src/core/optimization/cache/compile/cache_graph_lowering.hpp b/src/core/optimization/cache/compile/cache_graph_lowering.hpp index 76696aad..8300297d 100644 --- a/src/core/optimization/cache/compile/cache_graph_lowering.hpp +++ b/src/core/optimization/cache/compile/cache_graph_lowering.hpp @@ -39,10 +39,12 @@ struct CacheRunnerHooks { // Null on the host path. std::function()> substep_capture_probe; - // Substep-path tap-driven DiCache probe: runs the shallow - // prefix and returns delta_y/delta_x/gamma computed on-device from taps + - // persistent operands, via the TapRegistry. (probe_depth). Null on the host path. - std::function substep_probe; + // Substep-path tap-driven DiCache probe: runs the shallow prefix and returns + // delta_y/delta_x/gamma. The metrics are woven by cache.rel_l1 / cache.gamma_indicator + // operators (resolved from the passed registry); the model supplies the probe-history + // device operands (still DiCacheGpuState-owned in the shallow migration). + // (probe_depth, operators). Null on the host path. + std::function substep_probe; // Host-path substep variants (models with no device slot — Wan). Tap-driven like // the device ones, but the residual / probe tensors are read back to host. diff --git a/src/core/optimization/cache/ir/graph_extension.hpp b/src/core/optimization/cache/ir/graph_extension.hpp index 23988c24..5fe4c167 100644 --- a/src/core/optimization/cache/ir/graph_extension.hpp +++ b/src/core/optimization/cache/ir/graph_extension.hpp @@ -33,7 +33,11 @@ struct GraphExtension { // cache slot by the pass handoff (MagCache residual capture). // ReplaceStream : the woven node replaces the block-stack stream at reuse // (wired in the reuse step; not produced yet). - enum class Sink { CaptureToSlot, ReplaceStream }; + // Indicator : the woven node is a scalar decision metric (DiCache + // delta_y/gamma); pinned under output_name (a "cache_ind:" + // name) and read back to host by run_substep_pass. No stream + // substitution and no slot d2d. + enum class Sink { CaptureToSlot, ReplaceStream, Indicator }; const ICacheOperator* op = nullptr; // resolved operator (op->lower emits nodes) std::string op_id; // operator id, for diagnostics diff --git a/src/core/optimization/cache/operator/cache_operators.cpp b/src/core/optimization/cache/operator/cache_operators.cpp index 96e1b09e..0026be0c 100644 --- a/src/core/optimization/cache/operator/cache_operators.cpp +++ b/src/core/optimization/cache/operator/cache_operators.cpp @@ -1,6 +1,7 @@ #include "core/optimization/cache/operator/cache_operator_registry.hpp" #include +#include #include "ggml.h" @@ -263,6 +264,107 @@ class GammaBlendOperator final : public ICacheOperator { } }; +// Scalar rel-L1 decision metric: out = sum|cur - ref| / sum|ref|. inputs: [cur, ref]. +// Matches build_probe_metrics rel() (ggml_extend.hpp). Used for DiCache delta_y +// (probe vs prev_probe) and delta_x (before vs prev_input). Emits a [1] scalar node; +// only that scalar is read back (the on-device red line — the ref operand is a +// full block-stack tensor that never round-trips). +class RelL1Operator final : public ICacheOperator { +public: + CacheOperatorSchema schema() const override { return {"cache.rel_l1", 2, 2, 1, true}; } + + bool apply_host(const std::vector*>& inputs, + const CacheOperatorParams&, + std::vector>* outputs) const override { + if (inputs.size() != 2 || inputs[0] == nullptr || inputs[1] == nullptr || + outputs == nullptr || !same_numel(*inputs[0], *inputs[1])) { + return false; + } + const int64_t n = inputs[0]->numel(); + const float* cur = inputs[0]->data(); + const float* ref = inputs[1]->data(); + double num = 0.0, den = 0.0; + for (int64_t i = 0; i < n; ++i) { + num += std::fabs(static_cast(cur[i]) - static_cast(ref[i])); + den += std::fabs(static_cast(ref[i])); + } + sd::Tensor out({1}); + out.data()[0] = den > 0.0 ? static_cast(num / den) : 0.0f; + outputs->assign(1, std::move(out)); + return true; + } + + bool lower(GraphLoweringContext& ctx, + const std::vector& inputs, + const CacheOperatorParams&, + std::vector* outputs) const override { + if (inputs.size() != 2 || inputs[0] == nullptr || inputs[1] == nullptr || + ctx.ctx == nullptr || outputs == nullptr) { + return false; + } + ggml_tensor* num = ggml_sum(ctx.ctx, ggml_abs(ctx.ctx, ggml_sub(ctx.ctx, inputs[0], inputs[1]))); + ggml_tensor* den = ggml_sum(ctx.ctx, ggml_abs(ctx.ctx, inputs[1])); + outputs->assign(1, ggml_div(ctx.ctx, num, den)); + return true; + } +}; + +// DiCache gamma decision metric: cur_resid = cur_probe - before; then +// out = sum|cur_resid - prev2| / sum|prev1 - prev2|. inputs: [cur_probe, before, +// prev1, prev2] (prev1/prev2 = the 2-deep probe-residual ring). Matches the +// hardcoded gamma weave. Emits a [1] scalar; only it is read back. +class GammaIndicatorOperator final : public ICacheOperator { +public: + CacheOperatorSchema schema() const override { return {"cache.gamma_indicator", 4, 4, 1, true}; } + + bool apply_host(const std::vector*>& inputs, + const CacheOperatorParams&, + std::vector>* outputs) const override { + if (inputs.size() != 4 || inputs[0] == nullptr || inputs[1] == nullptr || + inputs[2] == nullptr || inputs[3] == nullptr || outputs == nullptr) { + return false; + } + const int64_t n = inputs[0]->numel(); + if (inputs[1]->numel() != n || inputs[2]->numel() != n || inputs[3]->numel() != n) { + return false; + } + const float* probe = inputs[0]->data(); + const float* before = inputs[1]->data(); + const float* prev1 = inputs[2]->data(); + const float* prev2 = inputs[3]->data(); + double num = 0.0, den = 0.0; + for (int64_t i = 0; i < n; ++i) { + const double cur_resid = static_cast(probe[i]) - static_cast(before[i]); + num += std::fabs(cur_resid - static_cast(prev2[i])); + den += std::fabs(static_cast(prev1[i]) - static_cast(prev2[i])); + } + sd::Tensor out({1}); + out.data()[0] = den > 0.0 ? static_cast(num / den) : 0.0f; + outputs->assign(1, std::move(out)); + return true; + } + + // cur_resid = sub(probe,before); num = sum|cur_resid - prev2|; + // den = sum|prev1 - prev2|; out = div(num,den). Byte-matches the hardcoded weave. + bool lower(GraphLoweringContext& ctx, + const std::vector& inputs, + const CacheOperatorParams&, + std::vector* outputs) const override { + if (inputs.size() != 4 || inputs[0] == nullptr || inputs[1] == nullptr || + inputs[2] == nullptr || inputs[3] == nullptr || ctx.ctx == nullptr || + outputs == nullptr) { + return false; + } + ggml_tensor* cur_resid = ggml_sub(ctx.ctx, inputs[0], inputs[1]); // probe - before + ggml_tensor* num = ggml_sum(ctx.ctx, ggml_abs(ctx.ctx, + ggml_sub(ctx.ctx, cur_resid, inputs[3]))); // |cur_resid - prev2| + ggml_tensor* den = ggml_sum(ctx.ctx, ggml_abs(ctx.ctx, + ggml_sub(ctx.ctx, inputs[2], inputs[3]))); // |prev1 - prev2| + outputs->assign(1, ggml_div(ctx.ctx, num, den)); + return true; + } +}; + } // namespace void register_builtin_cache_operators(CacheOperatorRegistry* registry) { @@ -276,6 +378,8 @@ void register_builtin_cache_operators(CacheOperatorRegistry* registry) { registry->register_operator(std::make_unique()); registry->register_operator(std::make_unique()); registry->register_operator(std::make_unique()); + registry->register_operator(std::make_unique()); + registry->register_operator(std::make_unique()); // cache.history_rotate is a state-manager action, not a tensor op; it has no // host lowering and is handled directly by CacheStateManager::rotate_history. } diff --git a/src/dit_models/models/flux.hpp b/src/dit_models/models/flux.hpp index 41410600..9900e88e 100644 --- a/src/dit_models/models/flux.hpp +++ b/src/dit_models/models/flux.hpp @@ -19,6 +19,7 @@ #include "edge-dit.h" #include "core/runtime/model_loader.h" #include "core/optimization/cache/ir/graph_extension.hpp" +#include "core/optimization/cache/operator/cache_operator_registry.hpp" #include "parallel/sp_parallel.hpp" #define FLUX_GRAPH_SIZE 10240 @@ -4997,7 +4998,8 @@ namespace Flux { bool increase_ref_index, int probe_depth, const void* branch_key, - bool delta_minus) { + bool delta_minus, + const edgedit::cache::CacheOperatorRegistry& operators) { const int m = std::max(1, probe_depth); edgedit::cache::TapRegistry reg; const auto probe_anchor = edgedit::cache::AnchorRef::block_out(m - 1); @@ -5005,22 +5007,51 @@ namespace Flux { reg.set_requested({before_anchor, probe_anchor}); reg.set_stop_after(m - 1); - edgedit::cache::TapRegistry::ProbeMetricOperands ops; + // Build the decision-metric extensions the runner weaves into the graph. + // A-shallow: the probe-history operands (prev_probe/prev_input and the + // 2-deep probe-residual ring) are still DiCacheGpuState-owned; the model + // backfills them into the extensions' extra_inputs. The reduction math + // lives in the cache operators (rel_l1 / gamma_indicator), not the runner. + std::vector exts; if (branch_key != nullptr) { auto it = dicache_gpu_states_.find(branch_key); if (it != dicache_gpu_states_.end() && it->second.has_probe_hist) { DiCacheGpuState& gpu = it->second; - ops.prev_probe = gpu.prev_probe; - ops.prev_input = gpu.prev_input; - ops.want_delta_x = delta_minus; + using edgedit::cache::GraphExtension; + // delta_y = rel_l1(probe, prev_probe) + GraphExtension dy; + dy.op = operators.find("cache.rel_l1"); + dy.op_id = "cache.rel_l1"; + dy.input_anchors = {probe_anchor}; + dy.extra_inputs = {gpu.prev_probe}; + dy.output_name = "cache_ind:delta_y"; + dy.sink = GraphExtension::Sink::Indicator; + if (dy.op != nullptr) exts.push_back(std::move(dy)); + if (delta_minus) { + // delta_x = rel_l1(before, prev_input) + GraphExtension dx; + dx.op = operators.find("cache.rel_l1"); + dx.op_id = "cache.rel_l1"; + dx.input_anchors = {before_anchor}; + dx.extra_inputs = {gpu.prev_input}; + dx.output_name = "cache_ind:delta_x"; + dx.sink = GraphExtension::Sink::Indicator; + if (dx.op != nullptr) exts.push_back(std::move(dx)); + } if (gpu.probe_resid_count >= 2) { - ops.probe_prev1 = gpu.probe_prev1; - ops.probe_prev2 = gpu.probe_prev2; - ops.want_gamma = true; + // gamma = gamma_indicator(probe, before, probe_prev1, probe_prev2) + GraphExtension gm; + gm.op = operators.find("cache.gamma_indicator"); + gm.op_id = "cache.gamma_indicator"; + gm.input_anchors = {probe_anchor, before_anchor}; + gm.extra_inputs = {gpu.probe_prev1, gpu.probe_prev2}; + gm.output_name = "cache_ind:gamma"; + gm.sink = GraphExtension::Sink::Indicator; + if (gm.op != nullptr) exts.push_back(std::move(gm)); } } } - reg.set_probe_metrics(probe_anchor, before_anchor, ops); + reg.set_extensions(std::move(exts)); auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, {}); diff --git a/src/dit_models/models/qwen_image.hpp b/src/dit_models/models/qwen_image.hpp index 3c8017e9..93e9de7e 100644 --- a/src/dit_models/models/qwen_image.hpp +++ b/src/dit_models/models/qwen_image.hpp @@ -21,6 +21,7 @@ #include "backend/ggml/ed_ggml_modulation_ext.hpp" #include "optimization/cache/compile/indicator_lowering.hpp" #include "optimization/cache/ir/graph_extension.hpp" +#include "optimization/cache/operator/cache_operator_registry.hpp" #endif namespace Qwen { @@ -3524,7 +3525,8 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* bool increase_ref_index, int probe_depth, const void* branch_key, - bool delta_minus) { + bool delta_minus, + const edgedit::cache::CacheOperatorRegistry& operators) { const int m = std::max(1, probe_depth); edgedit::cache::TapRegistry reg; const auto probe_anchor = edgedit::cache::AnchorRef::block_out(m - 1); @@ -3532,22 +3534,46 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* reg.set_requested({before_anchor, probe_anchor}); reg.set_stop_after(m - 1); // 0-based: return after block m-1 (m blocks run) - edgedit::cache::TapRegistry::ProbeMetricOperands ops; + // A-shallow: decision-metric reductions live in cache operators + // (rel_l1 / gamma_indicator); the probe-history device operands are still + // DiCacheGpuState-owned and backfilled into the extensions' extra_inputs. + std::vector exts; if (branch_key != nullptr) { auto it = dicache_gpu_states_.find(branch_key); if (it != dicache_gpu_states_.end() && it->second.has_probe_hist) { DiCacheGpuState& gpu = it->second; - ops.prev_probe = gpu.prev_probe; - ops.prev_input = gpu.prev_input; - ops.want_delta_x = delta_minus; + using edgedit::cache::GraphExtension; + GraphExtension dy; + dy.op = operators.find("cache.rel_l1"); + dy.op_id = "cache.rel_l1"; + dy.input_anchors = {probe_anchor}; + dy.extra_inputs = {gpu.prev_probe}; + dy.output_name = "cache_ind:delta_y"; + dy.sink = GraphExtension::Sink::Indicator; + if (dy.op != nullptr) exts.push_back(std::move(dy)); + if (delta_minus) { + GraphExtension dx; + dx.op = operators.find("cache.rel_l1"); + dx.op_id = "cache.rel_l1"; + dx.input_anchors = {before_anchor}; + dx.extra_inputs = {gpu.prev_input}; + dx.output_name = "cache_ind:delta_x"; + dx.sink = GraphExtension::Sink::Indicator; + if (dx.op != nullptr) exts.push_back(std::move(dx)); + } if (gpu.probe_resid_count >= 2) { - ops.probe_prev1 = gpu.probe_prev1; - ops.probe_prev2 = gpu.probe_prev2; - ops.want_gamma = true; + GraphExtension gm; + gm.op = operators.find("cache.gamma_indicator"); + gm.op_id = "cache.gamma_indicator"; + gm.input_anchors = {probe_anchor, before_anchor}; + gm.extra_inputs = {gpu.probe_prev1, gpu.probe_prev2}; + gm.output_name = "cache_ind:gamma"; + gm.sink = GraphExtension::Sink::Indicator; + if (gm.op != nullptr) exts.push_back(std::move(gm)); } } } - reg.set_probe_metrics(probe_anchor, before_anchor, ops); + reg.set_extensions(std::move(exts)); auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, ref_latents, increase_ref_index); diff --git a/src/dit_models/pipelines/flux_kontext_pipeline.cpp b/src/dit_models/pipelines/flux_kontext_pipeline.cpp index 5be6ef7b..38950f76 100644 --- a/src/dit_models/pipelines/flux_kontext_pipeline.cpp +++ b/src/dit_models/pipelines/flux_kontext_pipeline.cpp @@ -1051,11 +1051,11 @@ bool FluxKontextPipeline::generate_one_image(const ed_image_generation_params_t* } if (cache_runtime.granularity() == cache::CacheGranularity::Probe) { const bool delta_minus = cache_runtime.dicache_delta_minus(); - hooks.substep_probe = [&, branch_key, delta_minus](int depth) { + hooks.substep_probe = [&, branch_key, delta_minus](int depth, const cache::CacheOperatorRegistry& operators) { return flux_runner_->compute_substep_probe(n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, cond_in.c_vector, guidance, ref_latents, false, depth, branch_key, - delta_minus); + delta_minus, operators); }; if (Flux::FluxRunner::dicache_gpu_enabled()) { hooks.substep_inject_gpu = [&, branch_key](std::vector exts) { diff --git a/src/dit_models/pipelines/flux_pipeline.cpp b/src/dit_models/pipelines/flux_pipeline.cpp index 4ec14ed0..8a5877b2 100644 --- a/src/dit_models/pipelines/flux_pipeline.cpp +++ b/src/dit_models/pipelines/flux_pipeline.cpp @@ -965,11 +965,11 @@ bool FluxPipeline::generate_one_image(const ed_image_generation_params_t* params // Substep-path tap-driven probe: delta_y/gamma // on-device from taps + persistent operands, no CacheGraphScope. const bool delta_minus = cache_runtime.dicache_delta_minus(); - hooks.substep_probe = [&, branch_key, delta_minus](int depth) { + hooks.substep_probe = [&, branch_key, delta_minus](int depth, const cache::CacheOperatorRegistry& operators) { return flux_runner_->compute_substep_probe(n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, cond_in.c_vector, guidance, {}, false, depth, branch_key, - delta_minus); + delta_minus, operators); }; // Only wire the on-GPU inject when the model's GPU DiCache path // is active. With ED_DICACHE_GPU=0, leaving inject_gpu unset lets diff --git a/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp b/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp index 62081e32..a1e1e5ad 100644 --- a/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp +++ b/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp @@ -783,10 +783,10 @@ bool QwenImageEditPipeline::generate_one_image(const ed_image_generation_params_ } if (cache_runtime.granularity() == cache::CacheGranularity::Probe) { const bool delta_minus = cache_runtime.dicache_delta_minus(); - hooks.substep_probe = [&, cond_in, branch_key, delta_minus](int depth) { + hooks.substep_probe = [&, cond_in, branch_key, delta_minus](int depth, const cache::CacheOperatorRegistry& operators) { return diffusion_->compute_substep_probe(n_threads, x, timesteps, cond_in.c_crossattn, ref_latents, - true, depth, branch_key, delta_minus); + true, depth, branch_key, delta_minus, operators); }; if (Qwen::QwenImageRunner::dicache_gpu_enabled()) { hooks.substep_inject_gpu = [&, cond_in, branch_key](std::vector exts) { diff --git a/src/dit_models/pipelines/qwen_image_pipeline.cpp b/src/dit_models/pipelines/qwen_image_pipeline.cpp index 39ad9019..ea070987 100644 --- a/src/dit_models/pipelines/qwen_image_pipeline.cpp +++ b/src/dit_models/pipelines/qwen_image_pipeline.cpp @@ -732,10 +732,10 @@ bool QwenImagePipeline::generate_one_image(const ed_image_generation_params_t* p // Substep-path tap-driven probe: delta_y/gamma // computed on-device from taps + persistent operands, no scope. const bool delta_minus = cache_runtime.dicache_delta_minus(); - hooks.substep_probe = [&, cond_in, branch_key, delta_minus](int depth) { + hooks.substep_probe = [&, cond_in, branch_key, delta_minus](int depth, const cache::CacheOperatorRegistry& operators) { return diffusion_->compute_substep_probe(n_threads, x, timesteps, cond_in.c_crossattn, empty_ref_latents, - false, depth, branch_key, delta_minus); + false, depth, branch_key, delta_minus, operators); }; // Only wire on-GPU inject when the model's GPU DiCache path is // active; with ED_DICACHE_GPU=0 the lowering takes the declarative From 6c3a3dbeb43963ee0499030cf47e146d23d63201 Mon Sep 17 00:00:00 2001 From: Liwei Lan Date: Thu, 16 Jul 2026 01:28:59 +0800 Subject: [PATCH 8/8] refactor: move DiCache cross-step rings into CacheStateManager device slots --- .../cache/compile/cache_graph_lowering.cpp | 30 ++- .../cache/compile/cache_graph_lowering.hpp | 25 +- .../optimization/cache/ir/graph_extension.hpp | 27 +++ .../cache/policy/policy_common.hpp | 17 +- src/dit_models/models/flux.hpp | 221 +++++++----------- src/dit_models/models/qwen_image.hpp | 195 ++++++---------- .../pipelines/flux_kontext_pipeline.cpp | 27 ++- src/dit_models/pipelines/flux_pipeline.cpp | 39 ++-- .../pipelines/qwen_image_edit_pipeline.cpp | 26 +-- .../pipelines/qwen_image_pipeline.cpp | 33 ++- 10 files changed, 310 insertions(+), 330 deletions(-) diff --git a/src/core/optimization/cache/compile/cache_graph_lowering.cpp b/src/core/optimization/cache/compile/cache_graph_lowering.cpp index 56db5d19..bb293466 100644 --- a/src/core/optimization/cache/compile/cache_graph_lowering.cpp +++ b/src/core/optimization/cache/compile/cache_graph_lowering.cpp @@ -161,6 +161,30 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, policy.set_substep_input(hooks.input); policy.begin_substeps(step, condition_key); + // Bridge the DiCache multi-slot ring to the runner without exposing the state + // manager type. rotate/alloc/read/filled forward to CacheStateManager for this + // branch's condition_key. See DiCacheSlotBridge's depth-convention warning. + DiCacheSlotBridge dicache_bridge; + dicache_bridge.rotate = [&state, condition_key](int slot) { + state.rotate_history(condition_key, slot); + }; + dicache_bridge.alloc = [&state, condition_key](int slot, const std::vector& shape) -> void* { + return state.alloc_device_entry(condition_key, slot, shape); + }; + dicache_bridge.read = [&state, condition_key](int slot, int depth) -> void* { + CacheSlotHandle h = state.read_history(condition_key, slot, depth); + return h.valid ? h.buffer : nullptr; + }; + dicache_bridge.filled = [&state, condition_key](int slot) -> int { + // Derive fill count from read_history validity (no dedicated accessor): + // probe onto increasing depths until one is invalid. + int n = 0; + while (state.read_history(condition_key, slot, n).valid) { + ++n; + } + return n; + }; + const int slot0 = program.slots.empty() ? 0 : program.slots.front().id; const bool device_slot = state.has_device_store() && !program.slots.empty() && program.slots.front().device_backed && @@ -293,7 +317,7 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, // Tap-driven probe: delta_y/delta_x/gamma woven by cache operators // (rel_l1 / gamma_indicator), resolved from the registry; the model // supplies the probe-history device operands. - sd::DiffusionCacheResult pr = hooks.substep_probe(depth, operators); + sd::DiffusionCacheResult pr = hooks.substep_probe(depth, operators, dicache_bridge); result.indicators["delta_y"] = pr.delta_y; result.indicators["delta_x"] = pr.delta_x; result.indicators["gamma"] = pr.gamma; @@ -326,7 +350,7 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, exts.push_back(std::move(inj)); } if (!exts.empty()) { - y = hooks.substep_inject_gpu(std::move(exts)); + y = hooks.substep_inject_gpu(std::move(exts), dicache_bridge); } } if (y.empty()) { @@ -347,7 +371,7 @@ sd::Tensor CacheGraphLowering::execute_substeps(ICachePolicy& policy, // run_cache_pass + CacheGraphScope). Produces the step output. ---- if (plan.op.kind == SubstepOpKind::CaptureProbeSeed) { if (hooks.substep_capture_probe) { - y = hooks.substep_capture_probe(); + y = hooks.substep_capture_probe(dicache_bridge); if (!y.empty()) { CacheObservation obs; obs.kind = CacheObservation::Kind::Feature; diff --git a/src/core/optimization/cache/compile/cache_graph_lowering.hpp b/src/core/optimization/cache/compile/cache_graph_lowering.hpp index 8300297d..61e823cc 100644 --- a/src/core/optimization/cache/compile/cache_graph_lowering.hpp +++ b/src/core/optimization/cache/compile/cache_graph_lowering.hpp @@ -33,18 +33,18 @@ struct CacheRunnerHooks { std::function(std::vector)> substep_capture; // Substep-path tap-driven DiCache device seed capture: full forward whose - // post-readback refreshes the persistent DiCacheGpuState residual/probe rings - // device-to-device from the tap-woven nodes (no device slot — the ring is - // runner-owned, not a StateManager slot). Returns the step output. - // Null on the host path. - std::function()> substep_capture_probe; + // post-readback refreshes the cross-step residual/probe rings device-to-device + // from the tap-woven nodes. The ring lives in CacheStateManager device slots + // (face C); the bridge lets the runner drive rotate/alloc without depending on + // the state manager type. Returns the step output. Null on the host path. + std::function(DiCacheSlotBridge)> substep_capture_probe; // Substep-path tap-driven DiCache probe: runs the shallow prefix and returns // delta_y/delta_x/gamma. The metrics are woven by cache.rel_l1 / cache.gamma_indicator - // operators (resolved from the passed registry); the model supplies the probe-history - // device operands (still DiCacheGpuState-owned in the shallow migration). - // (probe_depth, operators). Null on the host path. - std::function substep_probe; + // operators (resolved from the passed registry); the probe-history device operands + // come from the CacheStateManager slots via the bridge. + // (probe_depth, operators, bridge). Null on the host path. + std::function substep_probe; // Host-path substep variants (models with no device slot — Wan). Tap-driven like // the device ones, but the residual / probe tensors are read back to host. @@ -63,9 +63,10 @@ struct CacheRunnerHooks { std::function(std::vector)> substep_inject_slot; // DiCache device reuse: the lowering hands over a ReplaceStream GraphExtension // (cache.gamma_blend, gamma baked into params.floats as a host constant); the - // model fills the residual-ring operands and build_stream_override weaves the - // blend. The runner never learns the math is a gamma extrapolation. - std::function(std::vector)> substep_inject_gpu; + // model reads the residual-ring operands from the CacheStateManager slots via + // the bridge and build_stream_override weaves the blend. The runner never learns + // the math is a gamma extrapolation. (extensions, bridge). + std::function(std::vector, DiCacheSlotBridge)> substep_inject_gpu; }; // Drives the policy's substep loop against the runner hooks, using the policy for diff --git a/src/core/optimization/cache/ir/graph_extension.hpp b/src/core/optimization/cache/ir/graph_extension.hpp index 5fe4c167..124e17eb 100644 --- a/src/core/optimization/cache/ir/graph_extension.hpp +++ b/src/core/optimization/cache/ir/graph_extension.hpp @@ -55,5 +55,32 @@ struct GraphExtension { std::function&)> alloc_slot; }; +// Callback bridge for DiCache's multi-slot ring, letting the runner drive the +// CacheStateManager device slots WITHOUT depending on the state manager type +// (same decoupling trick as GraphExtension::alloc_slot). The lowering binds these +// to state.rotate_history / alloc_device_entry / read_history for a fixed +// condition_key; the runner calls them by slot id. +// +// ⚠️ DEPTH CONVENTION (face C, mock-verified in /tmp/rotate_equiv.cpp): +// writeback is ROTATE-FIRST then write, so the NEWEST entry is at depth 0 +// (read(slot)=read_history(slot,0)) and the previous at depth 1. This DIVERGES +// from the host make_dicache_program, which reads of_slot_history(slot,1)+(,2). +// The device path never executes those host LOAD actions, so the divergence is +// benign — but if DiCache reuse/gamma ever reads the WRONG entry (off-by-one in +// the residual blend or a stale gamma), THIS is the first place to check: confirm +// rotate() is called BEFORE alloc()/write() in writeback, and that reads use +// depths 0/1 (newest/prev), not 1/2. +struct DiCacheSlotBridge { + // rotate the slot's ring so the next alloc targets a fresh newest entry. + std::function rotate; + // allocate/fetch the current newest ring entry at `shape`; returns device tensor. + std::function& /*shape*/)> alloc; + // read a history entry (depth 0 = newest after writeback); nullptr if unfilled. + std::function read; + // how many entries are filled in the slot's ring (availability gate). + std::function filled; + bool valid() const { return rotate && alloc && read && filled; } +}; + } // namespace cache } // namespace edgedit diff --git a/src/core/optimization/cache/policy/policy_common.hpp b/src/core/optimization/cache/policy/policy_common.hpp index f818070d..747d391d 100644 --- a/src/core/optimization/cache/policy/policy_common.hpp +++ b/src/core/optimization/cache/policy/policy_common.hpp @@ -286,7 +286,22 @@ inline CacheProgram make_dicache_program(const char* method, int block_segment_i int probe_depth) { CacheProgram program; program.method_name = method; - program.slots.push_back(make_slot(0, "block_stack_residual", 3)); + // Device-backed cross-step rings (face C). The runner's writeback drives these + // via the DiCacheSlotBridge; the host-path variants below still describe the + // logical STORE/REUSE for the (host) declarative interpreter. + // slot0: full residual ring (depth 2) — reuse gamma-blend reads newest+prev + // slot1: probe residual ring (depth 2) — gamma indicator reads newest+prev + // slot2: prev_probe snapshot (depth 1) — delta_y ref + // slot3: prev_input snapshot (depth 1) — delta_x ref + // ⚠️ DEVICE DEPTH CONVENTION DIVERGES from the of_slot_history(0,1)/(0,2) reads + // in the REUSE variant below: the device writeback is rotate-first so newest is + // at depth 0 (see DiCacheSlotBridge). The device path does NOT run the REUSE + // LOAD actions — it reads slots directly via the bridge at depths 0/1. If a + // future host-declarative DiCache device path is built, reconcile these. + program.slots.push_back(make_slot(0, "dicache_residual_ring", 2, /*device*/true)); + program.slots.push_back(make_slot(1, "dicache_probe_resid_ring", 2, /*device*/true)); + program.slots.push_back(make_slot(2, "dicache_prev_probe", 1, /*device*/true)); + program.slots.push_back(make_slot(3, "dicache_prev_input", 1, /*device*/true)); GraphVariantPlan full; full.id = kVariantFull; diff --git a/src/dit_models/models/flux.hpp b/src/dit_models/models/flux.hpp index 9900e88e..bf319cd3 100644 --- a/src/dit_models/models/flux.hpp +++ b/src/dit_models/models/flux.hpp @@ -4361,33 +4361,11 @@ namespace Flux { bool use_mask = false; sd::Tensor inject_feature_host_; // kept alive across cache inject build - // ---- Persistent cross-step GPU state for DiCache (Probe granularity) ---- - // Holds the last computed step's probe state, block input, probe residuals - // (2-deep) and full residuals (2-deep) as GPU tensors, per CFG branch, so - // the decision metric and residual reconstruction run on-device instead of - // reading ~50MB back to host each step. Allocated lazily at the live img - // shape; freed on reset / shape change. See dicache_perf notes. - struct DiCacheGpuState { - ggml_context* ctx = nullptr; - ggml_backend_buffer_t buffer = nullptr; - ggml_tensor* prev_probe = nullptr; - ggml_tensor* prev_input = nullptr; - ggml_tensor* probe_prev1 = nullptr; - ggml_tensor* probe_prev2 = nullptr; - ggml_tensor* resid_prev1 = nullptr; - ggml_tensor* resid_prev2 = nullptr; - std::vector shape; // [ne0, ne1, ne2] the tensors were built for - bool has_probe_hist = false; // prev_probe/prev_input valid - int probe_resid_count = 0; // 0/1/2 valid probe residuals - int resid_count = 0; // 0/1/2 valid full residuals - void free() { - if (buffer != nullptr) { ggml_backend_buffer_free(buffer); buffer = nullptr; } - if (ctx != nullptr) { ggml_free(ctx); ctx = nullptr; } - prev_probe = prev_input = probe_prev1 = probe_prev2 = resid_prev1 = resid_prev2 = nullptr; - shape.clear(); has_probe_hist = false; probe_resid_count = 0; resid_count = 0; - } - }; - std::unordered_map dicache_gpu_states_; + // ---- DiCache (Probe granularity) cross-step state ---- + // The residual/probe rings + prev-probe/input snapshots now live in + // CacheStateManager device slots (face C), reached via the DiCacheSlotBridge + // threaded through the substep hooks. The former per-branch DiCacheGpuState + // struct + dicache_gpu_states_ map were removed. FluxRunner(ggml_backend_t backend, bool offload_params_to_cpu, @@ -4805,45 +4783,10 @@ namespace Flux { } // ---- Cache seam passes (Feature/Probe policies). Reuse build_graph. ---- - // Ensure the per-branch persistent DiCache GPU state exists and matches the - // current block-stack shape [ne0, ne1, ne2]. (Re)allocates on first use or - // a shape change. Returns the state; its tensors are zero-initialized and - // marked invalid (has_probe_hist=false, counts=0) until the first capture - // fills them. - DiCacheGpuState& ensure_dicache_gpu_state(const void* branch_key, - int64_t ne0, int64_t ne1, int64_t ne2) { - DiCacheGpuState& s = dicache_gpu_states_[branch_key]; - const std::vector want = {ne0, ne1, ne2}; - if (s.buffer != nullptr && s.shape == want) { - return s; - } - s.free(); - s.ctx = new_cache_context(16); - auto mk = [&]() { - ggml_tensor* t = ggml_new_tensor_3d(s.ctx, GGML_TYPE_F32, ne0, ne1, ne2); - return t; - }; - s.prev_probe = mk(); - s.prev_input = mk(); - s.probe_prev1 = mk(); - s.probe_prev2 = mk(); - s.resid_prev1 = mk(); - s.resid_prev2 = mk(); - s.buffer = ggml_backend_alloc_ctx_tensors(s.ctx, runtime_backend); - GGML_ASSERT(s.buffer != nullptr); - s.shape = want; - s.has_probe_hist = false; - s.probe_resid_count = 0; - s.resid_count = 0; - return s; - } - - void reset_dicache_gpu_states() { - for (auto& kv : dicache_gpu_states_) { - kv.second.free(); - } - dicache_gpu_states_.clear(); - } + // DiCache's cross-step rings now live in CacheStateManager device slots + // (face C), driven via the DiCacheSlotBridge in the substep hooks. The old + // per-branch DiCacheGpuState + ensure/reset helpers were removed; lifecycle + // is now CacheStateManager::reset() (per generation). int dicache_probe_depth_ = 1; // set by the pipeline before capture/probe @@ -4908,19 +4851,23 @@ namespace Flux { const std::vector>& ref_latents, bool increase_ref_index, std::vector extensions, - const void* branch_key) { - auto it = dicache_gpu_states_.find(branch_key); - if (it == dicache_gpu_states_.end() || it->second.buffer == nullptr || - it->second.resid_count < 2 || extensions.empty()) { + const edgedit::cache::DiCacheSlotBridge& bridge) { + // Residual ring lives in CacheStateManager slot0 (face C). Need >=2 + // filled entries for the 2-tap gamma-blend. ⚠️ depth 0 = newest, 1 = + // prev (rotate-first writeback); NOT 1/2 (see DiCacheSlotBridge). + if (!bridge.valid() || extensions.empty() || bridge.filled(0) < 2) { return {}; // not enough history yet -> lowering falls back to full } - DiCacheGpuState& s = it->second; + ggml_tensor* resid_newest = static_cast(bridge.read(0, 0)); + ggml_tensor* resid_prev = static_cast(bridge.read(0, 1)); + if (resid_newest == nullptr || resid_prev == nullptr) { + return {}; + } edgedit::cache::TapRegistry reg; const int resume = flux_params.depth + flux_params.depth_single_blocks; - // Shallow migration: gamma is baked into the extension by the lowering; - // the model only supplies the residual-ring device operands. build_graph - // prepends x_before, so gamma_blend sees [x_before, resid_prev1, resid_prev2]. - extensions[0].extra_inputs = {s.resid_prev1, s.resid_prev2}; + // gamma baked into the extension by the lowering; build_graph prepends + // x_before, so gamma_blend sees [x_before, resid_newest(=prev1), resid_prev(=prev2)]. + extensions[0].extra_inputs = {resid_newest, resid_prev}; reg.set_extensions(std::move(extensions)); reg.set_override_region(0, resume); auto get_graph = [&]() -> ggml_cgraph* { @@ -4999,7 +4946,8 @@ namespace Flux { int probe_depth, const void* branch_key, bool delta_minus, - const edgedit::cache::CacheOperatorRegistry& operators) { + const edgedit::cache::CacheOperatorRegistry& operators, + const edgedit::cache::DiCacheSlotBridge& bridge) { const int m = std::max(1, probe_depth); edgedit::cache::TapRegistry reg; const auto probe_anchor = edgedit::cache::AnchorRef::block_out(m - 1); @@ -5007,44 +4955,50 @@ namespace Flux { reg.set_requested({before_anchor, probe_anchor}); reg.set_stop_after(m - 1); - // Build the decision-metric extensions the runner weaves into the graph. - // A-shallow: the probe-history operands (prev_probe/prev_input and the - // 2-deep probe-residual ring) are still DiCacheGpuState-owned; the model - // backfills them into the extensions' extra_inputs. The reduction math - // lives in the cache operators (rel_l1 / gamma_indicator), not the runner. + // Decision-metric extensions the runner weaves. The probe-history device + // operands now come from CacheStateManager slots via the bridge (face C): + // slot2 prev_probe (depth1), slot3 prev_input (depth1), + // slot1 probe-residual ring (depth0=newest, depth1=prev). + // ⚠️ depths 0/1, NOT 1/2 (rotate-first writeback; see DiCacheSlotBridge). + // history_ready mirrors the old has_probe_hist (prev_probe seeded). std::vector exts; - if (branch_key != nullptr) { - auto it = dicache_gpu_states_.find(branch_key); - if (it != dicache_gpu_states_.end() && it->second.has_probe_hist) { - DiCacheGpuState& gpu = it->second; - using edgedit::cache::GraphExtension; - // delta_y = rel_l1(probe, prev_probe) + const bool history_ready = bridge.valid() && bridge.filled(2) >= 1; + if (history_ready) { + using edgedit::cache::GraphExtension; + ggml_tensor* prev_probe = static_cast(bridge.read(2, 0)); + ggml_tensor* prev_input = static_cast(bridge.read(3, 0)); + // delta_y = rel_l1(probe, prev_probe) + if (prev_probe != nullptr) { GraphExtension dy; dy.op = operators.find("cache.rel_l1"); dy.op_id = "cache.rel_l1"; dy.input_anchors = {probe_anchor}; - dy.extra_inputs = {gpu.prev_probe}; + dy.extra_inputs = {prev_probe}; dy.output_name = "cache_ind:delta_y"; dy.sink = GraphExtension::Sink::Indicator; if (dy.op != nullptr) exts.push_back(std::move(dy)); - if (delta_minus) { - // delta_x = rel_l1(before, prev_input) - GraphExtension dx; - dx.op = operators.find("cache.rel_l1"); - dx.op_id = "cache.rel_l1"; - dx.input_anchors = {before_anchor}; - dx.extra_inputs = {gpu.prev_input}; - dx.output_name = "cache_ind:delta_x"; - dx.sink = GraphExtension::Sink::Indicator; - if (dx.op != nullptr) exts.push_back(std::move(dx)); - } - if (gpu.probe_resid_count >= 2) { - // gamma = gamma_indicator(probe, before, probe_prev1, probe_prev2) + } + if (delta_minus && prev_input != nullptr) { + // delta_x = rel_l1(before, prev_input) + GraphExtension dx; + dx.op = operators.find("cache.rel_l1"); + dx.op_id = "cache.rel_l1"; + dx.input_anchors = {before_anchor}; + dx.extra_inputs = {prev_input}; + dx.output_name = "cache_ind:delta_x"; + dx.sink = GraphExtension::Sink::Indicator; + if (dx.op != nullptr) exts.push_back(std::move(dx)); + } + if (bridge.filled(1) >= 2) { + // gamma = gamma_indicator(probe, before, probe_prev1, probe_prev2) + ggml_tensor* probe_prev1 = static_cast(bridge.read(1, 0)); + ggml_tensor* probe_prev2 = static_cast(bridge.read(1, 1)); + if (probe_prev1 != nullptr && probe_prev2 != nullptr) { GraphExtension gm; gm.op = operators.find("cache.gamma_indicator"); gm.op_id = "cache.gamma_indicator"; gm.input_anchors = {probe_anchor, before_anchor}; - gm.extra_inputs = {gpu.probe_prev1, gpu.probe_prev2}; + gm.extra_inputs = {probe_prev1, probe_prev2}; gm.output_name = "cache_ind:gamma"; gm.sink = GraphExtension::Sink::Indicator; if (gm.op != nullptr) exts.push_back(std::move(gm)); @@ -5070,42 +5024,47 @@ namespace Flux { return out; } - // ---- Substep-path tap-driven DiCache seed capture (device). Mirrors Qwen's - // compute_substep_capture_probe: a full forward whose post-readback refreshes - // the persistent DiCacheGpuState residual/probe rings device-to-device from the - // tap-woven residual/probe nodes. No device slot (the ring is runner-owned, - // not a StateManager slot). ---- - void writeback_from_taps(const void* branch_key) { + // ---- Substep-path tap-driven DiCache seed capture (device, face C). A full + // forward whose post-readback refreshes the cross-step residual/probe rings + // device-to-device from the tap-woven nodes, now into CacheStateManager + // device slots via the DiCacheSlotBridge (no more runner-owned DiCacheGpuState). ---- + // + // ⚠️ ORDER IS LOAD-BEARING (mock-verified /tmp/rotate_equiv.cpp): for the + // 2-deep rings we ROTATE FIRST, then alloc the (new) newest head, then d2d. + // That makes read(slot,0)=newest, read(slot,1)=prev — matching the legacy + // swap(prev1,prev2)+overwrite-prev1. If DiCache reuse/gamma ever reads a + // stale or off-by-one residual, verify this rotate-before-write order and + // the depth-0-is-newest convention FIRST. + void writeback_to_slots(const edgedit::cache::DiCacheSlotBridge& bridge, int probe_depth) { + if (!bridge.valid()) { + return; + } ggml_tensor* feat = get_cache_tensor_by_name("ed_cache_feature"); if (feat == nullptr) { return; } - DiCacheGpuState& s = - ensure_dicache_gpu_state(branch_key, feat->ne[0], feat->ne[1], feat->ne[2]); - if (s.buffer == nullptr) { - return; + std::vector shape; + const int nd = std::max(1, ggml_n_dims(feat)); + for (int i = 0; i < nd; ++i) { + shape.push_back(feat->ne[i]); } - const int m = std::max(1, dicache_probe_depth_); + const int m = std::max(1, probe_depth); const std::string probe_tap = "ed_tap:block_out[" + std::to_string(m - 1) + "]"; - // A silent copy failure (shape/name mismatch) would leave the ring stale - // and silently degrade the decision, so surface it loudly. - auto cp = [&](const char* name, ggml_tensor* dst) { - if (!copy_named_cache_tensor_to(name, dst)) { + auto cp = [&](const char* name, void* dst) { + ggml_tensor* d = static_cast(dst); + if (d == nullptr || !copy_named_cache_tensor_to(name, d)) { LOG_ERROR("dicache writeback copy failed: %s", name); } }; - // full residual ring: prev2 <- prev1, prev1 <- feature - std::swap(s.resid_prev1, s.resid_prev2); - cp("ed_cache_feature", s.resid_prev1); - s.resid_count = std::min(2, s.resid_count + 1); - // probe residual ring: prev2 <- prev1, prev1 <- (probe - before) - std::swap(s.probe_prev1, s.probe_prev2); - cp("ed_cache_probe_resid", s.probe_prev1); - s.probe_resid_count = std::min(2, s.probe_resid_count + 1); - // raw probe & input for the delta_y / delta_x decision metrics - cp(probe_tap.c_str(), s.prev_probe); - cp("ed_tap:model_in", s.prev_input); - s.has_probe_hist = true; + // slot0 full residual ring (depth2): rotate, then write newest head. + bridge.rotate(0); + cp("ed_cache_feature", bridge.alloc(0, shape)); + // slot1 probe residual ring (depth2): rotate, then write newest head. + bridge.rotate(1); + cp("ed_cache_probe_resid", bridge.alloc(1, shape)); + // slot2 prev_probe / slot3 prev_input (depth1 snapshots): no rotate. + cp(probe_tap.c_str(), bridge.alloc(2, shape)); + cp("ed_tap:model_in", bridge.alloc(3, shape)); } sd::Tensor compute_substep_capture_probe(int n_threads, @@ -5118,7 +5077,7 @@ namespace Flux { const std::vector>& ref_latents, bool increase_ref_index, int probe_depth, - const void* branch_key) { + const edgedit::cache::DiCacheSlotBridge& bridge) { const int m = std::max(1, probe_depth); edgedit::cache::TapRegistry reg; reg.set_requested({edgedit::cache::AnchorRef::model_in(), @@ -5126,7 +5085,7 @@ namespace Flux { edgedit::cache::AnchorRef::model_out()}); reg.set_capture_residual(true); // weave ed_cache_feature = ModelOut - ModelIn reg.set_capture_writeback(m); // weave ed_cache_probe_resid = block_out(m-1) - ModelIn - std::function handoff = [&]() { writeback_from_taps(branch_key); }; + std::function handoff = [&]() { writeback_to_slots(bridge, m); }; auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, {}); }; diff --git a/src/dit_models/models/qwen_image.hpp b/src/dit_models/models/qwen_image.hpp index 93e9de7e..9a113d44 100644 --- a/src/dit_models/models/qwen_image.hpp +++ b/src/dit_models/models/qwen_image.hpp @@ -3097,34 +3097,12 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* SDVersion version; sd::Tensor inject_feature_host_; // kept alive across cache inject build - // ---- Persistent cross-step GPU state for DiCache (Probe granularity) ---- - // Mirrors FluxRunner::DiCacheGpuState: holds the last computed step's probe - // state, block input, probe residuals (2-deep) and full residuals (2-deep) - // as GPU tensors, per CFG branch, so the decision metric and residual - // reconstruction run on-device instead of reading ~50MB back to host each - // step. Allocated lazily at the live img shape; freed on reset / shape - // change. See dicache_perf notes. - struct DiCacheGpuState { - ggml_context* ctx = nullptr; - ggml_backend_buffer_t buffer = nullptr; - ggml_tensor* prev_probe = nullptr; - ggml_tensor* prev_input = nullptr; - ggml_tensor* probe_prev1 = nullptr; - ggml_tensor* probe_prev2 = nullptr; - ggml_tensor* resid_prev1 = nullptr; - ggml_tensor* resid_prev2 = nullptr; - std::vector shape; // [ne0, ne1, ne2] the tensors were built for - bool has_probe_hist = false; // prev_probe/prev_input valid - int probe_resid_count = 0; // 0/1/2 valid probe residuals - int resid_count = 0; // 0/1/2 valid full residuals - void free() { - if (buffer != nullptr) { ggml_backend_buffer_free(buffer); buffer = nullptr; } - if (ctx != nullptr) { ggml_free(ctx); ctx = nullptr; } - prev_probe = prev_input = probe_prev1 = probe_prev2 = resid_prev1 = resid_prev2 = nullptr; - shape.clear(); has_probe_hist = false; probe_resid_count = 0; resid_count = 0; - } - }; - std::unordered_map dicache_gpu_states_; + // ---- DiCache (Probe granularity) cross-step state ---- + // The residual/probe rings + prev-probe/input snapshots now live in + // CacheStateManager device slots (face C), reached via the DiCacheSlotBridge + // threaded through the substep hooks. The former per-branch DiCacheGpuState + // struct + dicache_gpu_states_ map + ensure/reset helpers were removed; + // lifecycle is now CacheStateManager::reset() (per generation). int dicache_probe_depth_ = 1; // set by the pipeline before capture/probe // GPU DiCache is the default: on-device metric + reconstruction avoids the @@ -3138,41 +3116,6 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* return v[0] != '0'; } - // Ensure the per-branch persistent DiCache GPU state exists and matches the - // current block-stack shape [ne0, ne1, ne2]. (Re)allocates on first use or a - // shape change; tensors are marked invalid until the first capture fills them. - DiCacheGpuState& ensure_dicache_gpu_state(const void* branch_key, - int64_t ne0, int64_t ne1, int64_t ne2) { - DiCacheGpuState& s = dicache_gpu_states_[branch_key]; - const std::vector want = {ne0, ne1, ne2}; - if (s.buffer != nullptr && s.shape == want) { - return s; - } - s.free(); - s.ctx = new_cache_context(16); - auto mk = [&]() { return ggml_new_tensor_3d(s.ctx, GGML_TYPE_F32, ne0, ne1, ne2); }; - s.prev_probe = mk(); - s.prev_input = mk(); - s.probe_prev1 = mk(); - s.probe_prev2 = mk(); - s.resid_prev1 = mk(); - s.resid_prev2 = mk(); - s.buffer = ggml_backend_alloc_ctx_tensors(s.ctx, runtime_backend); - GGML_ASSERT(s.buffer != nullptr); - s.shape = want; - s.has_probe_hist = false; - s.probe_resid_count = 0; - s.resid_count = 0; - return s; - } - - void reset_dicache_gpu_states() { - for (auto& kv : dicache_gpu_states_) { - kv.second.free(); - } - dicache_gpu_states_.clear(); - } - QwenImageRunner(ggml_backend_t backend, bool offload_params_to_cpu, const String2TensorStorage& tensor_storage_map = {}, @@ -3381,16 +3324,20 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* const std::vector>& ref_latents, bool increase_ref_index, std::vector extensions, - const void* branch_key) { - auto it = dicache_gpu_states_.find(branch_key); - if (it == dicache_gpu_states_.end() || it->second.buffer == nullptr || - it->second.resid_count < 2 || extensions.empty()) { + const edgedit::cache::DiCacheSlotBridge& bridge) { + // Residual ring in CacheStateManager slot0 (face C). ⚠️ depth 0=newest, + // 1=prev (rotate-first writeback), NOT 1/2 (see DiCacheSlotBridge). + if (!bridge.valid() || extensions.empty() || bridge.filled(0) < 2) { return {}; // not enough history yet -> lowering falls back to full } - DiCacheGpuState& s = it->second; + ggml_tensor* resid_newest = static_cast(bridge.read(0, 0)); + ggml_tensor* resid_prev = static_cast(bridge.read(0, 1)); + if (resid_newest == nullptr || resid_prev == nullptr) { + return {}; + } edgedit::cache::TapRegistry reg; const int resume = qwen_image_params.num_layers; - extensions[0].extra_inputs = {s.resid_prev1, s.resid_prev2}; + extensions[0].extra_inputs = {resid_newest, resid_prev}; reg.set_extensions(std::move(extensions)); reg.set_override_region(0, resume); auto get_graph = [&]() -> ggml_cgraph* { @@ -3448,44 +3395,40 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* return std::move(pass.output); } - // ---- Substep-path tap-driven DiCache seed capture. A full - // forward whose post-readback refreshes the persistent DiCacheGpuState rings - // device-to-device from the tap-woven residual/probe nodes. No device slot is - // involved (DiCache's StateManager slot is host-backed/vestigial; the real - // cross-step state is the DiCacheGpuState ring). - void writeback_from_taps(const void* branch_key) { - // Shape source: the tap-woven residual (ModelOut - ModelIn). - // Allocate the rings lazily at that shape. + // ---- Substep-path tap-driven DiCache seed capture (device, face C). Refreshes + // the cross-step residual/probe rings device-to-device into CacheStateManager + // device slots via the DiCacheSlotBridge (no more runner-owned DiCacheGpuState). + // ⚠️ ORDER IS LOAD-BEARING (mock-verified /tmp/rotate_equiv.cpp): 2-deep rings + // ROTATE FIRST then write the newest head → read(slot,0)=newest, (,1)=prev, + // matching the legacy swap. If reuse/gamma reads stale/off-by-one, check the + // rotate-before-write order + depth-0-is-newest convention FIRST. ---- + void writeback_to_slots(const edgedit::cache::DiCacheSlotBridge& bridge, int probe_depth) { + if (!bridge.valid()) { + return; + } ggml_tensor* feat = get_cache_tensor_by_name("ed_cache_feature"); if (feat == nullptr) { return; } - DiCacheGpuState& s = - ensure_dicache_gpu_state(branch_key, feat->ne[0], feat->ne[1], feat->ne[2]); - if (s.buffer == nullptr) { - return; + std::vector shape; + const int nd = std::max(1, ggml_n_dims(feat)); + for (int i = 0; i < nd; ++i) { + shape.push_back(feat->ne[i]); } - const int m = std::max(1, dicache_probe_depth_); + const int m = std::max(1, probe_depth); const std::string probe_tap = "ed_tap:block_out[" + std::to_string(m - 1) + "]"; - // A silent copy failure (shape/name mismatch) would leave the ring stale - // and silently degrade the decision, so surface it loudly. - auto cp = [&](const char* name, ggml_tensor* dst) { - if (!copy_named_cache_tensor_to(name, dst)) { + auto cp = [&](const char* name, void* dst) { + ggml_tensor* d = static_cast(dst); + if (d == nullptr || !copy_named_cache_tensor_to(name, d)) { LOG_ERROR("dicache writeback copy failed: %s", name); } }; - // full residual ring: prev2 <- prev1, prev1 <- feature - std::swap(s.resid_prev1, s.resid_prev2); - cp("ed_cache_feature", s.resid_prev1); - s.resid_count = std::min(2, s.resid_count + 1); - // probe residual ring: prev2 <- prev1, prev1 <- (probe - before) - std::swap(s.probe_prev1, s.probe_prev2); - cp("ed_cache_probe_resid", s.probe_prev1); - s.probe_resid_count = std::min(2, s.probe_resid_count + 1); - // raw probe & input for the delta_y / delta_x decision metrics - cp(probe_tap.c_str(), s.prev_probe); - cp("ed_tap:model_in", s.prev_input); - s.has_probe_hist = true; + bridge.rotate(0); + cp("ed_cache_feature", bridge.alloc(0, shape)); // slot0 residual ring + bridge.rotate(1); + cp("ed_cache_probe_resid", bridge.alloc(1, shape)); // slot1 probe residual ring + cp(probe_tap.c_str(), bridge.alloc(2, shape)); // slot2 prev_probe (depth1) + cp("ed_tap:model_in", bridge.alloc(3, shape)); // slot3 prev_input (depth1) } sd::Tensor compute_substep_capture_probe(int n_threads, @@ -3495,7 +3438,7 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* const std::vector>& ref_latents, bool increase_ref_index, int probe_depth, - const void* branch_key) { + const edgedit::cache::DiCacheSlotBridge& bridge) { const int m = std::max(1, probe_depth); edgedit::cache::TapRegistry reg; // Full compute (no stop_after): block_out(m-1) is tapped in-loop anyway. @@ -3504,7 +3447,7 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* edgedit::cache::AnchorRef::model_out()}); reg.set_capture_residual(true); // weave ed_cache_feature = ModelOut - ModelIn reg.set_capture_writeback(m); // weave ed_cache_probe_resid = block_out(m-1) - ModelIn - std::function handoff = [&]() { writeback_from_taps(branch_key); }; + std::function handoff = [&]() { writeback_to_slots(bridge, m); }; auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, ref_latents, increase_ref_index); }; @@ -3526,7 +3469,8 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* int probe_depth, const void* branch_key, bool delta_minus, - const edgedit::cache::CacheOperatorRegistry& operators) { + const edgedit::cache::CacheOperatorRegistry& operators, + const edgedit::cache::DiCacheSlotBridge& bridge) { const int m = std::max(1, probe_depth); edgedit::cache::TapRegistry reg; const auto probe_anchor = edgedit::cache::AnchorRef::block_out(m - 1); @@ -3534,39 +3478,46 @@ static inline ggml_tensor* qwen_fused_attn_head_to_seq_recv_unpack(ggml_context* reg.set_requested({before_anchor, probe_anchor}); reg.set_stop_after(m - 1); // 0-based: return after block m-1 (m blocks run) - // A-shallow: decision-metric reductions live in cache operators - // (rel_l1 / gamma_indicator); the probe-history device operands are still - // DiCacheGpuState-owned and backfilled into the extensions' extra_inputs. + // Decision-metric reductions live in cache operators; the probe-history + // device operands come from CacheStateManager slots via the bridge (face C): + // slot2 prev_probe (depth1), slot3 prev_input (depth1), + // slot1 probe-residual ring (depth0=newest, depth1=prev). + // ⚠️ depths 0/1, NOT 1/2 (rotate-first writeback; see DiCacheSlotBridge). std::vector exts; - if (branch_key != nullptr) { - auto it = dicache_gpu_states_.find(branch_key); - if (it != dicache_gpu_states_.end() && it->second.has_probe_hist) { - DiCacheGpuState& gpu = it->second; - using edgedit::cache::GraphExtension; + const bool history_ready = bridge.valid() && bridge.filled(2) >= 1; + if (history_ready) { + using edgedit::cache::GraphExtension; + ggml_tensor* prev_probe = static_cast(bridge.read(2, 0)); + ggml_tensor* prev_input = static_cast(bridge.read(3, 0)); + if (prev_probe != nullptr) { GraphExtension dy; dy.op = operators.find("cache.rel_l1"); dy.op_id = "cache.rel_l1"; dy.input_anchors = {probe_anchor}; - dy.extra_inputs = {gpu.prev_probe}; + dy.extra_inputs = {prev_probe}; dy.output_name = "cache_ind:delta_y"; dy.sink = GraphExtension::Sink::Indicator; if (dy.op != nullptr) exts.push_back(std::move(dy)); - if (delta_minus) { - GraphExtension dx; - dx.op = operators.find("cache.rel_l1"); - dx.op_id = "cache.rel_l1"; - dx.input_anchors = {before_anchor}; - dx.extra_inputs = {gpu.prev_input}; - dx.output_name = "cache_ind:delta_x"; - dx.sink = GraphExtension::Sink::Indicator; - if (dx.op != nullptr) exts.push_back(std::move(dx)); - } - if (gpu.probe_resid_count >= 2) { + } + if (delta_minus && prev_input != nullptr) { + GraphExtension dx; + dx.op = operators.find("cache.rel_l1"); + dx.op_id = "cache.rel_l1"; + dx.input_anchors = {before_anchor}; + dx.extra_inputs = {prev_input}; + dx.output_name = "cache_ind:delta_x"; + dx.sink = GraphExtension::Sink::Indicator; + if (dx.op != nullptr) exts.push_back(std::move(dx)); + } + if (bridge.filled(1) >= 2) { + ggml_tensor* probe_prev1 = static_cast(bridge.read(1, 0)); + ggml_tensor* probe_prev2 = static_cast(bridge.read(1, 1)); + if (probe_prev1 != nullptr && probe_prev2 != nullptr) { GraphExtension gm; gm.op = operators.find("cache.gamma_indicator"); gm.op_id = "cache.gamma_indicator"; gm.input_anchors = {probe_anchor, before_anchor}; - gm.extra_inputs = {gpu.probe_prev1, gpu.probe_prev2}; + gm.extra_inputs = {probe_prev1, probe_prev2}; gm.output_name = "cache_ind:gamma"; gm.sink = GraphExtension::Sink::Indicator; if (gm.op != nullptr) exts.push_back(std::move(gm)); diff --git a/src/dit_models/pipelines/flux_kontext_pipeline.cpp b/src/dit_models/pipelines/flux_kontext_pipeline.cpp index 38950f76..c0ddd4c3 100644 --- a/src/dit_models/pipelines/flux_kontext_pipeline.cpp +++ b/src/dit_models/pipelines/flux_kontext_pipeline.cpp @@ -963,22 +963,19 @@ bool FluxKontextPipeline::generate_one_image(const ed_image_generation_params_t* parallel::cfg_parallel_available(runtime_->parallel_context()); const bool cache_seam_available = !cache_use_cfg_parallel && flux_runner_->feature_cache_available(); - // Wire the device store only when the on-GPU feature-reuse path is active - // (ED_FEATURE_CACHE_GPU on); with it off, leave the store null so a - // device_backed slot cleanly falls back to the host declarative path. + // Wire the device store when an on-GPU device path is active: MagCache + // feature-reuse (ED_FEATURE_CACHE_GPU) OR DiCache rings (ED_DICACHE_GPU, face C). cache::ICacheDeviceStore* cache_store = (cache_seam_available && flux_runner_ != nullptr && - Flux::FluxRunner::feature_gpu_enabled()) + (Flux::FluxRunner::feature_gpu_enabled() || Flux::FluxRunner::dicache_gpu_enabled())) ? flux_runner_->cache_device_store() : nullptr; const bool cache_enabled = cache_runtime.init(params->sample, version_, sigmas, cache_seam_available, cache_store, cache_use_cfg_parallel); - // GPU DiCache (ED_DICACHE_GPU): reset per-generation persistent state and set - // the probe depth the capture step uses to snapshot its probe residual. Read - // the resolved depth from the engine so it stays in sync with the policy config. + // GPU DiCache (ED_DICACHE_GPU): set the probe depth for the capture snapshot. + // Per-generation ring state is owned + freed by CacheStateManager::reset() (face C). if (cache_enabled && flux_runner_ != nullptr) { - flux_runner_->reset_dicache_gpu_states(); flux_runner_->dicache_probe_depth_ = cache_runtime.dicache_probe_depth(); } const int64_t sample_start_ms = ggml_time_ms(); @@ -1051,23 +1048,25 @@ bool FluxKontextPipeline::generate_one_image(const ed_image_generation_params_t* } if (cache_runtime.granularity() == cache::CacheGranularity::Probe) { const bool delta_minus = cache_runtime.dicache_delta_minus(); - hooks.substep_probe = [&, branch_key, delta_minus](int depth, const cache::CacheOperatorRegistry& operators) { + hooks.substep_probe = [&, branch_key, delta_minus](int depth, const cache::CacheOperatorRegistry& operators, + const cache::DiCacheSlotBridge& bridge) { return flux_runner_->compute_substep_probe(n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, cond_in.c_vector, guidance, ref_latents, false, depth, branch_key, - delta_minus, operators); + delta_minus, operators, bridge); }; if (Flux::FluxRunner::dicache_gpu_enabled()) { - hooks.substep_inject_gpu = [&, branch_key](std::vector exts) { + hooks.substep_inject_gpu = [&](std::vector exts, + const cache::DiCacheSlotBridge& bridge) { return flux_runner_->compute_substep_inject_gpu(n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, cond_in.c_vector, - guidance, ref_latents, false, std::move(exts), branch_key); + guidance, ref_latents, false, std::move(exts), bridge); }; const int probe_depth = cache_runtime.dicache_probe_depth(); - hooks.substep_capture_probe = [&, branch_key, probe_depth]() { + hooks.substep_capture_probe = [&, probe_depth](const cache::DiCacheSlotBridge& bridge) { return flux_runner_->compute_substep_capture_probe( n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, cond_in.c_vector, - guidance, ref_latents, false, probe_depth, branch_key); + guidance, ref_latents, false, probe_depth, bridge); }; } } diff --git a/src/dit_models/pipelines/flux_pipeline.cpp b/src/dit_models/pipelines/flux_pipeline.cpp index 8a5877b2..f18930e6 100644 --- a/src/dit_models/pipelines/flux_pipeline.cpp +++ b/src/dit_models/pipelines/flux_pipeline.cpp @@ -864,23 +864,26 @@ bool FluxPipeline::generate_one_image(const ed_image_generation_params_t* params !uncond.empty() && parallel::cfg_parallel_available(runtime_->parallel_context()); const bool cache_seam_available = !cfg_parallel_for_cache && flux_runner_->feature_cache_available(); - // Wire the device store only when the on-GPU feature-reuse path is active - // (ED_FEATURE_CACHE_GPU on); with it off, leave the store null so a - // device_backed slot cleanly falls back to the host declarative path. + // Wire the device store when an on-GPU device path is active: MagCache + // feature-reuse (ED_FEATURE_CACHE_GPU) OR DiCache's residual/probe rings + // (ED_DICACHE_GPU, face C — the rings are now CacheStateManager device slots). + // With both off, leave the store null so device_backed slots fall back to the + // host declarative path. The store is harmless if a run doesn't touch it (slots + // allocate lazily). cache::ICacheDeviceStore* cache_store = (cache_seam_available && flux_runner_ != nullptr && - Flux::FluxRunner::feature_gpu_enabled()) + (Flux::FluxRunner::feature_gpu_enabled() || Flux::FluxRunner::dicache_gpu_enabled())) ? flux_runner_->cache_device_store() : nullptr; const bool cache_enabled = cache_runtime.init(params->sample, version_, sigmas, cache_seam_available, cache_store, cfg_parallel_for_cache); - // GPU DiCache (ED_DICACHE_GPU): reset per-generation persistent state and set - // the probe depth the capture step uses to snapshot its probe residual. Read - // the resolved depth from the engine so it stays in sync with the policy's - // config (the reference default is 1, NOT DBCache's cache_Fn_compute_blocks). + // GPU DiCache: set the probe depth the capture step uses to snapshot its probe + // residual. Read the resolved depth from the engine so it stays in sync with the + // policy's config (the reference default is 1, NOT DBCache's cache_Fn_compute_blocks). + // Per-generation ring state is now owned + freed by CacheStateManager::reset() + // (face C); no more reset_dicache_gpu_states() here. if (cache_enabled && flux_runner_ != nullptr) { - flux_runner_->reset_dicache_gpu_states(); flux_runner_->dicache_probe_depth_ = cache_runtime.dicache_probe_depth(); } const int64_t sample_start_ms = ggml_time_ms(); @@ -965,11 +968,12 @@ bool FluxPipeline::generate_one_image(const ed_image_generation_params_t* params // Substep-path tap-driven probe: delta_y/gamma // on-device from taps + persistent operands, no CacheGraphScope. const bool delta_minus = cache_runtime.dicache_delta_minus(); - hooks.substep_probe = [&, branch_key, delta_minus](int depth, const cache::CacheOperatorRegistry& operators) { + hooks.substep_probe = [&, delta_minus](int depth, const cache::CacheOperatorRegistry& operators, + const cache::DiCacheSlotBridge& bridge) { return flux_runner_->compute_substep_probe(n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, cond_in.c_vector, guidance, {}, false, depth, branch_key, - delta_minus, operators); + delta_minus, operators, bridge); }; // Only wire the on-GPU inject when the model's GPU DiCache path // is active. With ED_DICACHE_GPU=0, leaving inject_gpu unset lets @@ -977,19 +981,20 @@ bool FluxPipeline::generate_one_image(const ed_image_generation_params_t* params // ring blend), instead of the legacy on-device reconstruction. if (Flux::FluxRunner::dicache_gpu_enabled()) { // Substep-path tap-driven device inject (DiCache gamma-blend). - hooks.substep_inject_gpu = [&, branch_key](std::vector exts) { + hooks.substep_inject_gpu = [&](std::vector exts, + const cache::DiCacheSlotBridge& bridge) { return flux_runner_->compute_substep_inject_gpu(n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, cond_in.c_vector, - guidance, {}, false, std::move(exts), branch_key); + guidance, {}, false, std::move(exts), bridge); }; // Substep-path tap-driven seed capture: full forward that - // refreshes the DiCacheGpuState rings device-to-device (replaces - // the legacy compute_capture / run_cache_pass path). + // refreshes the DiCache rings (CacheStateManager device slots, + // face C) device-to-device via the bridge. const int probe_depth = cache_runtime.dicache_probe_depth(); - hooks.substep_capture_probe = [&, branch_key, probe_depth]() { + hooks.substep_capture_probe = [&, probe_depth](const cache::DiCacheSlotBridge& bridge) { return flux_runner_->compute_substep_capture_probe( n_threads, noised_input, timesteps, cond_in.c_crossattn, {}, cond_in.c_vector, - guidance, {}, false, probe_depth, branch_key); + guidance, {}, false, probe_depth, bridge); }; } } diff --git a/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp b/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp index a1e1e5ad..19e9e639 100644 --- a/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp +++ b/src/dit_models/pipelines/qwen_image_edit_pipeline.cpp @@ -701,21 +701,19 @@ bool QwenImageEditPipeline::generate_one_image(const ed_image_generation_params_ parallel::cfg_parallel_available(runtime_->parallel_context()); const bool cache_seam_available = !cache_use_cfg_parallel && diffusion_->feature_cache_available(); - // Wire the device store only when the on-GPU feature-reuse path is active - // (ED_FEATURE_CACHE_GPU on); with it off, leave the store null so a - // device_backed slot cleanly falls back to the host declarative path. + // Wire the device store when an on-GPU device path is active: MagCache + // feature-reuse (ED_FEATURE_CACHE_GPU) OR DiCache rings (ED_DICACHE_GPU, face C). cache::ICacheDeviceStore* cache_store = (cache_seam_available && diffusion_ != nullptr && - Qwen::QwenImageRunner::feature_gpu_enabled()) + (Qwen::QwenImageRunner::feature_gpu_enabled() || Qwen::QwenImageRunner::dicache_gpu_enabled())) ? diffusion_->cache_device_store() : nullptr; const bool cache_enabled = cache_runtime.init(params->sample, version_, sigmas, cache_seam_available, cache_store, cache_use_cfg_parallel); - // GPU DiCache (ED_DICACHE_GPU): reset per-generation persistent state and set - // the probe depth the capture step uses, kept in sync with the policy config. + // GPU DiCache (ED_DICACHE_GPU): set the probe depth for the capture snapshot. + // Per-generation ring state is owned + freed by CacheStateManager::reset() (face C). if (cache_enabled && diffusion_ != nullptr) { - diffusion_->reset_dicache_gpu_states(); diffusion_->dicache_probe_depth_ = cache_runtime.dicache_probe_depth(); } const int64_t sample_start_ms = ggml_time_ms(); @@ -783,21 +781,23 @@ bool QwenImageEditPipeline::generate_one_image(const ed_image_generation_params_ } if (cache_runtime.granularity() == cache::CacheGranularity::Probe) { const bool delta_minus = cache_runtime.dicache_delta_minus(); - hooks.substep_probe = [&, cond_in, branch_key, delta_minus](int depth, const cache::CacheOperatorRegistry& operators) { + hooks.substep_probe = [&, cond_in, branch_key, delta_minus](int depth, const cache::CacheOperatorRegistry& operators, + const cache::DiCacheSlotBridge& bridge) { return diffusion_->compute_substep_probe(n_threads, x, timesteps, cond_in.c_crossattn, ref_latents, - true, depth, branch_key, delta_minus, operators); + true, depth, branch_key, delta_minus, operators, bridge); }; if (Qwen::QwenImageRunner::dicache_gpu_enabled()) { - hooks.substep_inject_gpu = [&, cond_in, branch_key](std::vector exts) { + hooks.substep_inject_gpu = [&, cond_in](std::vector exts, + const cache::DiCacheSlotBridge& bridge) { return diffusion_->compute_substep_inject_gpu(n_threads, x, timesteps, cond_in.c_crossattn, - ref_latents, true, std::move(exts), branch_key); + ref_latents, true, std::move(exts), bridge); }; const int probe_depth = cache_runtime.dicache_probe_depth(); - hooks.substep_capture_probe = [&, cond_in, branch_key, probe_depth]() { + hooks.substep_capture_probe = [&, cond_in, probe_depth](const cache::DiCacheSlotBridge& bridge) { return diffusion_->compute_substep_capture_probe( n_threads, x, timesteps, cond_in.c_crossattn, ref_latents, - true, probe_depth, branch_key); + true, probe_depth, bridge); }; } } diff --git a/src/dit_models/pipelines/qwen_image_pipeline.cpp b/src/dit_models/pipelines/qwen_image_pipeline.cpp index ea070987..b2e23759 100644 --- a/src/dit_models/pipelines/qwen_image_pipeline.cpp +++ b/src/dit_models/pipelines/qwen_image_pipeline.cpp @@ -634,22 +634,20 @@ bool QwenImagePipeline::generate_one_image(const ed_image_generation_params_t* p parallel::cfg_parallel_available(runtime_->parallel_context()); const bool cache_seam_available = !cache_use_cfg_parallel && diffusion_->feature_cache_available(); - // Wire the device store only when the on-GPU feature-reuse path is active - // (ED_FEATURE_CACHE_GPU on); with it off, leave the store null so a - // device_backed slot cleanly falls back to the host declarative path. + // Wire the device store when an on-GPU device path is active: MagCache + // feature-reuse (ED_FEATURE_CACHE_GPU) OR DiCache rings (ED_DICACHE_GPU, face C). + // With both off, device_backed slots fall back to the host declarative path. cache::ICacheDeviceStore* cache_store = (cache_seam_available && diffusion_ != nullptr && - Qwen::QwenImageRunner::feature_gpu_enabled()) + (Qwen::QwenImageRunner::feature_gpu_enabled() || Qwen::QwenImageRunner::dicache_gpu_enabled())) ? diffusion_->cache_device_store() : nullptr; const bool cache_enabled = cache_runtime.init(params->sample, version_, sigmas, cache_seam_available, cache_store, cache_use_cfg_parallel); - // GPU DiCache (ED_DICACHE_GPU): reset per-generation persistent state and set - // the probe depth the capture step uses to snapshot its probe residual. Read - // the resolved depth from the engine so it stays in sync with the policy config. + // GPU DiCache (ED_DICACHE_GPU): set the probe depth for the capture snapshot. + // Per-generation ring state is owned + freed by CacheStateManager::reset() (face C). if (cache_enabled && diffusion_ != nullptr) { - diffusion_->reset_dicache_gpu_states(); diffusion_->dicache_probe_depth_ = cache_runtime.dicache_probe_depth(); } const int64_t sample_start_ms = ggml_time_ms(); @@ -732,28 +730,29 @@ bool QwenImagePipeline::generate_one_image(const ed_image_generation_params_t* p // Substep-path tap-driven probe: delta_y/gamma // computed on-device from taps + persistent operands, no scope. const bool delta_minus = cache_runtime.dicache_delta_minus(); - hooks.substep_probe = [&, cond_in, branch_key, delta_minus](int depth, const cache::CacheOperatorRegistry& operators) { + hooks.substep_probe = [&, cond_in, branch_key, delta_minus](int depth, const cache::CacheOperatorRegistry& operators, + const cache::DiCacheSlotBridge& bridge) { return diffusion_->compute_substep_probe(n_threads, x, timesteps, cond_in.c_crossattn, empty_ref_latents, - false, depth, branch_key, delta_minus, operators); + false, depth, branch_key, delta_minus, operators, bridge); }; // Only wire on-GPU inject when the model's GPU DiCache path is // active; with ED_DICACHE_GPU=0 the lowering takes the declarative // host probe path instead (residual-ring blend). if (Qwen::QwenImageRunner::dicache_gpu_enabled()) { // Substep-path tap-driven device inject (DiCache gamma-blend). - hooks.substep_inject_gpu = [&, cond_in, branch_key](std::vector exts) { + hooks.substep_inject_gpu = [&, cond_in](std::vector exts, + const cache::DiCacheSlotBridge& bridge) { return diffusion_->compute_substep_inject_gpu(n_threads, x, timesteps, cond_in.c_crossattn, - empty_ref_latents, false, std::move(exts), branch_key); + empty_ref_latents, false, std::move(exts), bridge); }; - // Substep-path tap-driven seed capture: full forward that - // refreshes the DiCacheGpuState rings device-to-device (replaces - // the legacy compute_capture / run_cache_pass path). + // Substep-path tap-driven seed capture: refreshes the DiCache + // rings (CacheStateManager device slots, face C) via the bridge. const int probe_depth = cache_runtime.dicache_probe_depth(); - hooks.substep_capture_probe = [&, cond_in, branch_key, probe_depth]() { + hooks.substep_capture_probe = [&, cond_in, probe_depth](const cache::DiCacheSlotBridge& bridge) { return diffusion_->compute_substep_capture_probe( n_threads, x, timesteps, cond_in.c_crossattn, empty_ref_latents, - false, probe_depth, branch_key); + false, probe_depth, bridge); }; } }