ci: automatically close PRs from organization-owned forks

.github/workflows/close-organization-fork-prs.yml

@@ -0,0 +1,61 @@
+name: Close PRs from organization forks
+
+on:
+  pull_request_target:
+    types: [opened, reopened]
+
+permissions:
+  pull-requests: write
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
+  cancel-in-progress: false
+
+jobs:
+  close-organization-fork-pr:
+    if: >-
+      github.event.pull_request.head.repo.owner.type == 'Organization' &&
+      github.event.pull_request.head.repo.id != github.event.pull_request.base.repo.id
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    steps:
+      - name: Explain the contribution policy and close the PR
+        uses: actions/github-script@v9
+        with:
+          script: |
+            const { data: pr } = await github.rest.pulls.get({
+              ...context.repo,
+              pull_number: context.issue.number,
+            });
+            const headRepo = pr.head.repo;
+            if (pr.state !== 'open' || !headRepo ||
+                headRepo.id === pr.base.repo.id || headRepo.owner.type !== 'Organization') {
+              return;
+            }
+
+            const marker = '<!-- organization-fork-policy -->';
+            const comments = await github.paginate(github.rest.issues.listComments, {
+              ...context.repo,
+              issue_number: pr.number,
+              per_page: 100,
+            });
+            const alreadyExplained = comments.some(comment =>
+              comment.user?.login === 'github-actions[bot]' && comment.body?.includes(marker));
+            if (!alreadyExplained) {
+              await github.rest.issues.createComment({
+                ...context.repo,
+                issue_number: pr.number,
+                body: [
+                  marker,
+                  'This repository requires contributions from forks to use a personal fork with **Allow edits from maintainers** enabled.',
+                  'GitHub does not support this option for organization-owned forks, so this PR is being closed automatically.',
+                  'Please open a new PR from a fork in your personal GitHub account and enable **Allow edits from maintainers** so maintainers can help update the branch.',
+                  'See [the GitHub documentation](https://docs.github.com/en/pull-requests/how-tos/work-with-forks/allowing-changes-to-a-pull-request-branch-created-from-a-fork).',
+                ].join('\n\n'),
+              });
+            }
+            await github.rest.pulls.update({
+              ...context.repo,
+              pull_number: pr.number,
+              state: 'closed',
+            });

CONTRIBUTING.md

@@ -12,6 +12,10 @@ If you want to update a third-party dependency, please open an issue first inste
 
 ## Pull Requests
 
+When contributing from a fork, use a fork under your personal GitHub account and enable **Allow edits from maintainers**. This lets maintainers make follow-up fixes directly on the PR branch.
+
+PRs from organization-owned forks are automatically closed when opened or reopened because GitHub does not support this maintainer-edit option for those forks. Submit the changes from a personal fork instead. See [GitHub's documentation](https://docs.github.com/en/pull-requests/how-tos/work-with-forks/allowing-changes-to-a-pull-request-branch-created-from-a-fork).
+
 Keep each PR focused on one clear change. Large or overly complex PRs are harder to review and may not be merged.
 
 Do not include test code or test scripts in commits or PRs. Keep them local and report verification results in the PR description.

docs/qwen_image_2.1.md

@@ -40,6 +40,14 @@ Pass the reference image with `-r` and describe the edit in `-p`. Vision weights
 
 For multiple reference images, repeat `-r` in the desired order, for example `-r first.png -r second.png`.
 
+### Prefix cache
+
+By default, the first denoising call for each fixed condition saves the text and reference-image keys and values from every transformer layer. Later calls only compute the target-image tokens. Positive and negative conditions use separate caches, which are released when sampling ends.
+
+The cache uses FP32 on all attention backends. For the default 32-layer model, a prefix of 4096 tokens takes about 4 GiB per condition, in addition to weights and working buffers. The runner accounts for the cache when checking the memory budget. If a cached execution runs out of memory, it releases the prefix caches, disables caching for the rest of that sampling run, and retries the full sequence once. Per-step conditioning extensions currently use the full-sequence path.
+
+Disable this optimization with `--model-args qwen_image_2_1_prefix_cache=false`. It reuses step-independent activations; numerical results can still differ slightly because the matrix sizes change.
+
 ### Alpha channel
 
 This model supports alpha channel output. As the model determines whether to output a regular image or with transparency through the prompt, according to [official recommendation](https://github.com/QwenLM/Qwen-Image-2.1#transparent-image-generation-rgba), use the following prompt format for better results:

examples/common/common.cpp

@@ -518,7 +518,7 @@ ArgOptions SDContextParams::get_options() {
         {"",
          "--model-args",
          "extra model args, key=value list. Supports chroma_use_dit_mask, chroma_use_t5_mask, "
-         "chroma_t5_mask_pad, qwen_image_zero_cond_t",
+         "chroma_t5_mask_pad, qwen_image_zero_cond_t, qwen_image_2_1_prefix_cache",
          (int)',',
          &model_args},
         {"",

src/core/ggml_runner.cpp

@@ -644,6 +644,10 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
     std::optional<sd::Tensor<float>> output;
     try {
         output = execute_graph(graph, n_threads, no_return, read_outputs);
+    } catch (const std::bad_alloc&) {
+        last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
+        LOG_ERROR("%s graph allocation failed", get_desc().c_str());
+        return std::nullopt;
     } catch (const std::exception& error) {
         last_compute_status_ = GGML_STATUS_FAILED;
         LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(),
@@ -964,10 +968,16 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
         }
         LOG_DEBUG("%s executing segment %zu/%zu: %s", get_desc().c_str(),
                   index + 1, plan.segments.size(), segment.group_name.c_str());
-        if (!execute_segment(segment_graph, n_threads) ||
-            !cache_.capture(segment_graph) ||
-            !cut_cache_.capture(graph, segment, get_desc().c_str())) {
-            return fail_segment("execution or output caching");
+        if (!execute_segment(segment_graph, n_threads)) {
+            return fail_segment("execution");
+        }
+        auto cache_status = cache_.capture(segment_graph);
+        if (cache_status == GGML_STATUS_SUCCESS) {
+            cache_status = cut_cache_.capture(graph, segment, get_desc().c_str());
+        }
+        if (cache_status != GGML_STATUS_SUCCESS) {
+            last_compute_status_ = cache_status;
+            return fail_segment("output caching");
         }
         sync_runtime_residency();
         if (last) {

src/core/runner_cache.cpp

@@ -26,10 +26,13 @@ namespace sd {
 
     std::unique_ptr<CachedTensor> CachedTensor::copy(ggml_backend_t backend,
                                                      const std::string& name,
-                                                     ggml_tensor* source) {
+                                                     ggml_tensor* source,
+                                                     ggml_status& status) {
+        status = GGML_STATUS_FAILED;
         if (ggml_graph_cut::tensor_buffer(source) == nullptr) {
             return nullptr;
         }
+        status         = GGML_STATUS_ALLOC_FAILED;
         auto entry     = std::make_unique<CachedTensor>();
         entry->context = ggml_init({2 * ggml_tensor_overhead(), nullptr, true});
         if (entry->context == nullptr) {
@@ -50,6 +53,7 @@ namespace sd {
         } else {
             ggml_backend_tensor_copy(source, entry->tensor);
         }
+        status = GGML_STATUS_SUCCESS;
         return entry;
     }
 
@@ -106,9 +110,9 @@ namespace sd {
         return pending > SIZE_MAX - committed ? SIZE_MAX : committed + pending;
     }
 
-    bool RunnerCache::capture(ggml_cgraph* graph) {
+    ggml_status RunnerCache::capture(ggml_cgraph* graph) {
         if (outputs_.empty()) {
-            return true;
+            return GGML_STATUS_SUCCESS;
         }
         const auto tensors = cache_graph_tensors(graph);
         for (const auto& output : outputs_) {
@@ -116,14 +120,15 @@ namespace sd {
                 continue;
             }
             GGML_ASSERT(ggml_is_contiguous(output.second));
-            auto entry = CachedTensor::copy(backend_, output.first, output.second);
+            ggml_status status;
+            auto entry = CachedTensor::copy(backend_, output.first, output.second, status);
             if (entry == nullptr) {
-                return false;
+                return status;
             }
             pending_[output.first] = std::move(entry);
         }
         ggml_backend_synchronize(backend_);
-        return true;
+        return GGML_STATUS_SUCCESS;
     }
 
     void RunnerCache::graph_end(bool success) {
@@ -180,9 +185,9 @@ namespace sd {
         }
     }
 
-    bool GraphCutTensorCache::capture(ggml_cgraph* graph,
-                                      const ggml_graph_cut::Segment& segment,
-                                      const char* log_desc) {
+    ggml_status GraphCutTensorCache::capture(ggml_cgraph* graph,
+                                             const ggml_graph_cut::Segment& segment,
+                                             const char* log_desc) {
         size_t copied_bytes = 0;
         size_t copied_count = 0;
         for (int index : segment.output_node_indices) {
@@ -191,10 +196,11 @@ namespace sd {
                 !segment.future_cut_names.count(output->name)) {
                 continue;
             }
-            auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output));
+            ggml_status status;
+            auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output), status);
             if (entry == nullptr) {
                 LOG_ERROR("%s failed to capture graph cut tensor: %s", log_desc, output->name);
-                return false;
+                return status;
             }
             const size_t size = ggml_backend_buffer_get_size(entry->buffer);
             copied_bytes      = size > SIZE_MAX - copied_bytes ? SIZE_MAX : copied_bytes + size;
@@ -206,6 +212,6 @@ namespace sd {
             LOG_DEBUG("%s graph cut cache added %6.2f MB (%zu tensors)",
                       log_desc, copied_bytes / (1024.f * 1024.f), copied_count);
         }
-        return true;
+        return GGML_STATUS_SUCCESS;
     }
 }

src/core/runner_cache.h

@@ -20,7 +20,8 @@ namespace sd {
         ~CachedTensor();
         static std::unique_ptr<CachedTensor> copy(ggml_backend_t backend,
                                                   const std::string& name,
-                                                  ggml_tensor* source);
+                                                  ggml_tensor* source,
+                                                  ggml_status& status);
     };
     using CachedTensors = std::map<std::string, std::unique_ptr<CachedTensor>>;
 
@@ -41,7 +42,8 @@ namespace sd {
         const std::map<std::string, ggml_tensor*>& outputs() const { return outputs_; }
         size_t pending_bytes(ggml_cgraph* graph) const;
         size_t resident_bytes(ggml_backend_dev_t device) const;
-        bool capture(ggml_cgraph* graph);
+        bool empty() const { return committed_.empty(); }
+        ggml_status capture(ggml_cgraph* graph);
         void graph_end(bool success);
         void clear();
     };
@@ -57,7 +59,7 @@ namespace sd {
         size_t resident_bytes(ggml_backend_dev_t device) const;
         size_t estimate_output_bytes(ggml_cgraph* graph,
                                      const ggml_graph_cut::Segment& segment) const;
-        bool capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc);
+        ggml_status capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc);
         void prune(const std::unordered_set<std::string>& keep_names);
         void clear() { tensors_.clear(); }
     };

src/model/diffusion/minimax_h3.hpp

@@ -264,6 +264,9 @@ namespace MiniMaxH3 {
             for (int64_t i = 0; i < num_layers; ++i) {
                 auto block = std::dynamic_pointer_cast<TokenRefinerBlock>(blocks["blocks." + std::to_string(i)]);
                 x          = block->forward(ctx, x);
+                sd::ggml_graph_cut::mark_graph_cut(x,
+                                                   "minimax_h3.token_refiner.blocks." + std::to_string(i),
+                                                   "hidden_states");
             }
             return std::dynamic_pointer_cast<RMSNorm>(blocks["final_norm"])->forward(ctx, x);
         }
@@ -527,7 +530,11 @@ namespace MiniMaxH3 {
             GGML_ASSERT(context->ne[0] == config.text_dim);
             auto condition_proj = std::dynamic_pointer_cast<Linear>(blocks["condition_proj"]);
             auto token_refiner  = std::dynamic_pointer_cast<TokenRefiner>(blocks["token_refiner"]);
-            return token_refiner->forward(ctx, condition_proj->forward(ctx, context));
+            auto projected      = condition_proj->forward(ctx, context);
+            sd::ggml_graph_cut::mark_graph_cut(projected,
+                                               "minimax_h3.condition_proj",
+                                               "hidden_states");
+            return token_refiner->forward(ctx, projected);
         }
 
         ggml_tensor* time_embedding(GGMLRunnerContext* ctx,

src/model/diffusion/model.hpp

@@ -71,6 +71,8 @@ struct AnimaDiffusionExtra {
 
 struct QwenImage21DiffusionExtra {
     const sd::Tensor<int32_t>* image_slots = nullptr;
+    // Nonzero IDs identify immutable prefix inputs within one sampling run.
+    uint64_t prefix_id = 0;
 };
 
 struct WanDiffusionExtra {

src/model/diffusion/qwen_image_2_1.hpp

@@ -121,6 +121,18 @@ namespace Qwen {
         }
     };
 
+    struct QwenImage21PrefixCache {
+        enum class Mode {
+            NONE,
+            STORE,
+            REUSE
+        };
+        Mode mode = Mode::NONE;
+        std::string name;
+        std::string cut_group;
+        int64_t prefix_length = 0;
+    };
+
     class QwenImage21ZeroCenterRMSNorm : public RMSNorm {
     public:
         using RMSNorm::RMSNorm;
@@ -160,27 +172,49 @@ namespace Qwen {
             }
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::vector<QwenImage21Segment>& segments, const std::vector<ggml_tensor*>& masks) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::vector<QwenImage21Segment>& segments, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
             int64_t heads = x->ne[0] / dim_head;
             auto project  = [&](const char* name) {
                 auto h = std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, x);
                 return ggml_reshape_4d(ctx->ggml_ctx, h, dim_head, heads, x->ne[1], x->ne[2]);
             };
-            auto q              = project("to_q");
-            auto k              = project("to_k");
-            auto v              = project("to_v");
-            q                   = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"])->forward(ctx, q);
-            k                   = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"])->forward(ctx, k);
-            q                   = Rope::apply_rope(ctx->ggml_ctx, q, pe);
-            k                   = Rope::apply_rope(ctx->ggml_ctx, k, pe);
+            auto q = project("to_q");
+            auto k = project("to_k");
+            auto v = project("to_v");
+            q      = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"])->forward(ctx, q);
+            k      = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"])->forward(ctx, k);
+            q      = Rope::apply_rope(ctx->ggml_ctx, q, pe);
+            k      = Rope::apply_rope(ctx->ggml_ctx, k, pe);
+            if (cache.mode == QwenImage21PrefixCache::Mode::STORE) {
+                auto persist = [&](ggml_tensor* tensor, int axis, const char* name) {
+                    auto part = ggml_ext_slice(ctx->ggml_ctx, tensor, axis, 0, cache.prefix_length);
+                    auto copy = ggml_new_tensor(ctx->ggml_ctx, GGML_TYPE_F32, 4, part->ne);
+                    copy      = ggml_cpy(ctx->ggml_ctx, part, copy);
+                    // Keep the copy in this layer's segment so graph cuts do not
+                    // retain or recompute the full-sequence K/V in the final segment.
+                    sd::ggml_graph_cut::mark_graph_cut(copy, cache.cut_group, name);
+                    ctx->persist_cache_tensor(cache.name + "." + name, copy);
+                };
+                persist(k, 1, "k");
+                persist(v, 2, "v");
+            }
             ggml_tensor* result = nullptr;
-            for (size_t i = 0; i < segments.size(); ++i) {
-                const auto& segment = segments[i];
-                auto sq             = ggml_ext_slice(ctx->ggml_ctx, q, 1, segment.start, segment.end);
-                auto sk             = ggml_ext_slice(ctx->ggml_ctx, k, 1, 0, segment.end);
-                auto sv             = ggml_ext_slice(ctx->ggml_ctx, v, 2, 0, segment.end);
-                auto out            = ggml_ext_attention_ext(ctx, sq, sk, sv, heads, masks[i], true, ctx->flash_attn_enabled);
-                result              = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1);
+            if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) {
+                auto prefix_k = ctx->load_cache_tensor(cache.name + ".k");
+                auto prefix_v = ctx->load_cache_tensor(cache.name + ".v");
+                GGML_ASSERT(prefix_k != nullptr && prefix_v != nullptr);
+                k      = ggml_concat(ctx->ggml_ctx, prefix_k, k, 1);
+                v      = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2);
+                result = ggml_ext_attention_ext(ctx, q, k, v, heads, nullptr, true, ctx->flash_attn_enabled);
+            } else {
+                for (size_t i = 0; i < segments.size(); ++i) {
+                    const auto& segment = segments[i];
+                    auto sq             = ggml_ext_slice(ctx->ggml_ctx, q, 1, segment.start, segment.end);
+                    auto sk             = ggml_ext_slice(ctx->ggml_ctx, k, 1, 0, segment.end);
+                    auto sv             = ggml_ext_slice(ctx->ggml_ctx, v, 2, 0, segment.end);
+                    auto out            = ggml_ext_attention_ext(ctx, sq, sk, sv, heads, masks[i], true, ctx->flash_attn_enabled);
+                    result              = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1);
+                }
             }
             auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
             if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
@@ -219,13 +253,14 @@ namespace Qwen {
             return ggml_concat(ctx, prefix, target, 1);
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector<ggml_tensor*>& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks) {
-            auto h = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm1"])->forward(ctx, x);
-            h      = modulate(ctx->ggml_ctx, h, modulation[0], layout.prefix_length);
-            h      = std::dynamic_pointer_cast<QwenImage21Attention>(blocks["attn"])->forward(ctx, h, pe, layout.segments, masks);
-            x      = ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[1], layout.prefix_length, true));
-            h      = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm2"])->forward(ctx, x);
-            h      = modulate(ctx->ggml_ctx, h, modulation[2], layout.prefix_length);
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector<ggml_tensor*>& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
+            const int64_t prefix_length = cache.mode == QwenImage21PrefixCache::Mode::REUSE ? 0 : layout.prefix_length;
+            auto h                      = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm1"])->forward(ctx, x);
+            h                           = modulate(ctx->ggml_ctx, h, modulation[0], prefix_length);
+            h                           = std::dynamic_pointer_cast<QwenImage21Attention>(blocks["attn"])->forward(ctx, h, pe, layout.segments, masks, cache);
+            x                           = ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[1], prefix_length, true));
+            h                           = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm2"])->forward(ctx, x);
+            h                           = modulate(ctx->ggml_ctx, h, modulation[2], prefix_length);
             ggml_tensor* gate;
             auto fused = blocks.find("img_mlp.gate_up");
             if (fused != blocks.end()) {
@@ -239,7 +274,7 @@ namespace Qwen {
             }
             h = ggml_mul(ctx->ggml_ctx, h, ggml_silu(ctx->ggml_ctx, gate));
             h = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.out"])->forward(ctx, h);
-            return ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[3], layout.prefix_length, true));
+            return ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[3], prefix_length, true));
         }
     };
 
@@ -261,7 +296,7 @@ namespace Qwen {
             }
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, const std::vector<ggml_tensor*>& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, const std::vector<ggml_tensor*>& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
             auto time = ggml_concat(ctx->ggml_ctx, timestep, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), 0);
             // Runtime flow timesteps already use the [0, 1000] scale.
             time               = ggml_ext_timestep_embedding(ctx->ggml_ctx, time, 256, 10000, 1.f);
@@ -269,27 +304,37 @@ namespace Qwen {
             time               = ggml_silu(ctx->ggml_ctx, time);
             auto modulation    = std::dynamic_pointer_cast<Linear>(blocks["modulation.1"])->forward(ctx, time);
             auto mod           = ggml_ext_chunk(ctx->ggml_ctx, modulation, 4, 0);
-            auto text          = std::dynamic_pointer_cast<QwenImage21TextProjection>(blocks["txt_in"])->forward(ctx, context);
             auto img_in        = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
             ggml_tensor* joint = nullptr;
-            for (const auto& segment : layout.segments) {
-                ggml_tensor* h;
-                if (segment.image_index < 0) {
-                    h = ggml_ext_slice(ctx->ggml_ctx, text, 1, segment.context_start,
-                                       segment.context_start + segment.end - segment.start);
-                } else {
-                    auto image = segment.image_index == static_cast<int>(refs.size()) ? x : refs[segment.image_index];
-                    h          = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, image, 1, 1));
+            if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) {
+                joint = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, x, 1, 1));
+            } else {
+                auto text = std::dynamic_pointer_cast<QwenImage21TextProjection>(blocks["txt_in"])->forward(ctx, context);
+                for (const auto& segment : layout.segments) {
+                    ggml_tensor* h;
+                    if (segment.image_index < 0) {
+                        h = ggml_ext_slice(ctx->ggml_ctx, text, 1, segment.context_start,
+                                           segment.context_start + segment.end - segment.start);
+                    } else {
+                        auto image = segment.image_index == static_cast<int>(refs.size()) ? x : refs[segment.image_index];
+                        h          = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, image, 1, 1));
+                    }
+                    joint = joint == nullptr ? h : ggml_concat(ctx->ggml_ctx, joint, h, 1);
                 }
-                joint = joint == nullptr ? h : ggml_concat(ctx->ggml_ctx, joint, h, 1);
             }
             sd::ggml_graph_cut::mark_graph_cut(joint, "qwen_image_2_1.prelude", "joint");
             for (int i = 0; i < config.num_layers; ++i) {
-                auto block = std::dynamic_pointer_cast<QwenImage21TransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]);
-                joint      = block->forward(ctx, joint, mod, pe, layout, masks);
-                sd::ggml_graph_cut::mark_graph_cut(joint, "qwen_image_2_1.transformer_blocks." + std::to_string(i), "joint");
+                const std::string layer = "transformer_blocks." + std::to_string(i);
+                auto layer_cache        = cache;
+                layer_cache.name        = cache.name + "." + std::to_string(i);
+                layer_cache.cut_group   = "qwen_image_2_1." + layer;
+                auto block              = std::dynamic_pointer_cast<QwenImage21TransformerBlock>(blocks[layer]);
+                joint                   = block->forward(ctx, joint, mod, pe, layout, masks, layer_cache);
+                sd::ggml_graph_cut::mark_graph_cut(joint, layer_cache.cut_group, "joint");
+            }
+            if (cache.mode != QwenImage21PrefixCache::Mode::REUSE) {
+                joint = ggml_ext_slice(ctx->ggml_ctx, joint, 1, layout.prefix_length, joint->ne[1]);
             }
-            joint      = ggml_ext_slice(ctx->ggml_ctx, joint, 1, layout.prefix_length, joint->ne[1]);
             auto scale = std::dynamic_pointer_cast<Linear>(blocks["norm_out.linear"])->forward(ctx, ggml_ext_chunk(ctx->ggml_ctx, time, 2, 1)[0]);
             joint      = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out.norm"])->forward(ctx, joint);
             joint      = ggml_mul(ctx->ggml_ctx, joint, ggml_scale_bias(ctx->ggml_ctx, scale, 1.f, 1.f));
@@ -303,11 +348,18 @@ namespace Qwen {
         QwenImage21Model model;
         std::vector<float> pe_data;
         std::vector<sd::Tensor<float>> mask_data;
+        bool prefix_cache_enabled  = true;
+        bool prefix_cache_disabled = false;
 
-        QwenImage21Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        QwenImage21Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr, const char* model_args = nullptr)
             : DiffusionModelRunner(backend, prefix, weight_manager),
               config(QwenImage21Config::detect_from_weights(weights, prefix)),
               model(config) {
+            for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
+                if (key == "qwen_image_2_1_prefix_cache" && !parse_strict_bool(value, prefix_cache_enabled)) {
+                    LOG_WARN("ignoring invalid Qwen Image 2.1 model arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            }
             model.init(params_ctx, weights, prefix);
         }
 
@@ -317,6 +369,22 @@ namespace Qwen {
             model.get_param_tensors(tensors, prefix);
         }
 
+        bool has_prefix_cache(const QwenImage21PrefixCache& cache) {
+            for (int i = 0; i < config.num_layers; ++i) {
+                const auto name = cache.name + "." + std::to_string(i);
+                auto k          = get_cache_tensor_by_name(name + ".k");
+                auto v          = get_cache_tensor_by_name(name + ".v");
+                if (k == nullptr || v == nullptr || k->type != GGML_TYPE_F32 || v->type != GGML_TYPE_F32 ||
+                    k->ne[0] != config.head_dim || k->ne[1] != cache.prefix_length ||
+                    k->ne[2] != config.hidden_size / config.head_dim || k->ne[3] != 1 ||
+                    v->ne[0] != config.head_dim || v->ne[1] != config.hidden_size / config.head_dim ||
+                    v->ne[2] != cache.prefix_length || v->ne[3] != 1) {
+                    return false;
+                }
+            }
+            return true;
+        }
+
         sd::Tensor<float> compute(int n_threads, const DiffusionParams& inputs) override {
             const auto& x       = tensor_or_empty(inputs.x);
             const auto& context = tensor_or_empty(inputs.context);
@@ -345,38 +413,75 @@ namespace Qwen {
                 LOG_ERROR("%s", error.what());
                 return {};
             }
-            pe_data = Rope::embed_nd(layout.positions, 1, 10000.f, config.axes_dim);
-            mask_data.clear();
-            for (const auto& segment : layout.segments) {
-                sd::Tensor<float> mask;
-                if (segment.image_index < 0) {
-                    mask = sd::Tensor<float>::zeros({segment.end, segment.end - segment.start});
-                    for (int64_t q = segment.start; q < segment.end; ++q) {
-                        for (int64_t k = q + 1; k < segment.end; ++k) {
-                            mask[k + segment.end * (q - segment.start)] = -INFINITY;
+            if (!runner_started()) {
+                prefix_cache_disabled = false;
+            }
+            QwenImage21PrefixCache cache;
+            if (prefix_cache_enabled && !prefix_cache_disabled && extra != nullptr && extra->prefix_id != 0 && layout.prefix_length > 0) {
+                cache.name          = "qwen_image_2_1.prefix." + std::to_string(extra->prefix_id);
+                cache.prefix_length = layout.prefix_length;
+                cache.mode          = has_prefix_cache(cache) ? QwenImage21PrefixCache::Mode::REUSE : QwenImage21PrefixCache::Mode::STORE;
+            }
+            auto run = [&](const QwenImage21PrefixCache& active_cache) {
+                const bool cached         = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE;
+                const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0);
+                pe_data                   = Rope::embed_nd(std::vector<std::vector<float>>(first_position, layout.positions.end()), 1, 10000.f, config.axes_dim);
+                mask_data.clear();
+                if (!cached) {
+                    for (const auto& segment : layout.segments) {
+                        sd::Tensor<float> mask;
+                        if (segment.image_index < 0) {
+                            mask = sd::Tensor<float>::zeros({segment.end, segment.end - segment.start});
+                            for (int64_t q = segment.start; q < segment.end; ++q) {
+                                for (int64_t k = q + 1; k < segment.end; ++k) {
+                                    mask[k + segment.end * (q - segment.start)] = -INFINITY;
+                                }
+                            }
                         }
+                        mask_data.push_back(std::move(mask));
                     }
                 }
-                mask_data.push_back(std::move(mask));
-            }
-            auto build = [&]() {
-                auto graph = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE * 2);
-                auto pe    = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, layout.positions.size());
-                set_backend_tensor_data(pe, pe_data.data());
-                std::vector<ggml_tensor*> masks, ref_inputs;
-                for (const auto& mask : mask_data) {
-                    masks.push_back(mask.empty() ? nullptr : make_input(mask));
-                }
-                for (const auto& ref : refs) {
-                    ref_inputs.push_back(make_input(ref));
-                }
-                auto ctx = get_context();
-                auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), make_input(context),
-                                         ref_inputs, pe, layout, masks);
-                ggml_build_forward_expand(graph, out);
-                return graph;
+                auto build = [&]() {
+                    auto graph = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE * 2);
+                    auto pe    = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2,
+                                                    layout.positions.size() - (cached ? layout.prefix_length : 0));
+                    set_backend_tensor_data(pe, pe_data.data());
+                    std::vector<ggml_tensor*> masks, ref_inputs;
+                    for (const auto& mask : mask_data) {
+                        masks.push_back(mask.empty() ? nullptr : make_input(mask));
+                    }
+                    if (!cached) {
+                        for (const auto& ref : refs) {
+                            ref_inputs.push_back(make_input(ref));
+                        }
+                    }
+                    auto ctx = get_context();
+                    auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), cached ? nullptr : make_input(context),
+                                             ref_inputs, pe, layout, masks, active_cache);
+                    ggml_build_forward_expand(graph, out);
+                    return graph;
+                };
+                return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), x.dim());
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), x.dim());
+            auto result = run(cache);
+            if (result.empty() && last_compute_status() == GGML_STATUS_ALLOC_FAILED &&
+                (cache.mode != QwenImage21PrefixCache::Mode::NONE || !cache_.empty())) {
+                // The failed graph has ended before persistent inputs are released.
+                free_cache_ctx_and_buffer();
+                prefix_cache_disabled = true;
+                LOG_WARN("Qwen Image 2.1: insufficient memory for prefix caching; retrying without it for this sampling run");
+                return run(QwenImage21PrefixCache{});
+            }
+            if (!result.empty() && cache.mode == QwenImage21PrefixCache::Mode::STORE) {
+                if (!has_prefix_cache(cache)) {
+                    free_cache_ctx_and_buffer();
+                    prefix_cache_disabled = true;
+                    LOG_WARN("Qwen Image 2.1: incomplete prefix cache; disabling it for this sampling run");
+                } else {
+                    LOG_DEBUG("Qwen Image 2.1: cached prefix %" PRIu64 " (%" PRId64 " tokens)", extra->prefix_id, layout.prefix_length);
+                }
+            }
+            return result;
         }
     };
 }

src/model/vae/wan_vae.hpp

@@ -82,6 +82,20 @@ namespace WAN {
             }
 
             x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
+            if (w->ne[2] == 1 && x->ne[2] == 1 && x->ne[3] == in_channels) {
+                // One frame through a one-frame-deep kernel is a 2D conv; backends without
+                // im2col_3d (Metal) otherwise fall back to a much slower direct conv_3d.
+                if (!ggml_is_contiguous(x)) {
+                    x = ggml_cont(ctx->ggml_ctx, x);
+                }
+                ggml_tensor* x2 = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], x->ne[1], in_channels, 1);
+                ggml_tensor* w2 = ggml_reshape_4d(ctx->ggml_ctx, w, w->ne[0], w->ne[1], in_channels, out_channels);
+                x2              = ggml_ext_conv_2d(ctx->ggml_ctx, x2, w2, b,
+                                                   std::get<2>(stride), std::get<1>(stride), 0, 0,
+                                                   std::get<2>(dilation), std::get<1>(dilation),
+                                                   ctx->conv2d_direct_enabled);
+                return ggml_reshape_4d(ctx->ggml_ctx, x2, x2->ne[0], x2->ne[1], 1, out_channels);
+            }
             return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
                                     std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
                                     0, 0, 0,

src/pipeline/diffusion_engine.cpp

@@ -7,6 +7,7 @@
 #include <list>
 #include <mutex>
 #include <set>
+#include <tuple>
 #include <type_traits>
 #include <unordered_set>
 #include <utility>
@@ -2255,6 +2256,15 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
     };
     RunnerEndOnExit sample_diffusion_runner_end{work_diffusion_model.get()};
 
+    // These inputs are immutable for this sampling run. Extensions may replace or
+    // modify them per step, so those paths need an explicit stability contract first.
+    const bool cache_qwen_prefix = version == VERSION_QWEN_IMAGE_2_1 &&
+                                   std::none_of(generation_extensions.begin(), generation_extensions.end(),
+                                                [](const auto& extension) { return extension->is_enabled(); });
+    using QwenPrefixInputs = std::tuple<const sd::Tensor<float>*, const sd::Tensor<int32_t>*,
+                                        const std::vector<sd::Tensor<float>>*>;
+    std::vector<QwenPrefixInputs> qwen_prefix_inputs;
+
     RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr};
 
     const bool apply_denoise_mask = !denoise_mask.empty() &&
@@ -2524,6 +2534,18 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
                 extension->before_diffusion(diffusion_params, step);
             }
 
+            if (cache_qwen_prefix) {
+                auto* extra = std::get_if<QwenImage21DiffusionExtra>(&diffusion_params.extra);
+                if (extra != nullptr) {
+                    auto key         = std::make_tuple(diffusion_params.context, extra->image_slots,
+                                               diffusion_params.ref_image_params.pass_to_dit ? diffusion_params.ref_latents : nullptr);
+                    auto entry       = std::find(qwen_prefix_inputs.begin(), qwen_prefix_inputs.end(), key);
+                    extra->prefix_id = static_cast<uint64_t>(entry - qwen_prefix_inputs.begin()) + 1;
+                    if (entry == qwen_prefix_inputs.end()) {
+                        qwen_prefix_inputs.push_back(key);
+                    }
+                }
+            }
             auto output_opt = work_diffusion_model->compute(n_threads, diffusion_params);
             if (output_opt.empty()) {
                 LOG_ERROR("diffusion model compute failed");

src/pipeline/model_builders.cpp

@@ -291,7 +291,8 @@ namespace sd::model_builders {
                 result.diffusion = std::make_shared<Qwen::QwenImage21Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                              tensor_storage_map,
                                                                              "model.diffusion_model",
-                                                                             weight_manager);
+                                                                             weight_manager,
+                                                                             sd_ctx_params->model_args);
             } else {
                 result.diffusion = std::make_shared<Qwen::QwenImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                            tensor_storage_map,