refactor: centralize circular RoPE and extend image model support (#2039)

.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/backend.md

@@ -161,6 +161,9 @@ resident allocations. Vulkan reports exceeding total memory are rejected because
 its heap-budget subtraction can underflow. Other backends use the cap instead of
 treating such reports as zero free memory. Failed checks log the reported free and
 total memory alongside tracked weight and runtime allocations.
+With `--mmap`, device-backed mappings count toward these budgets at their full
+mapped-file size, once per device buffer even when multiple parameter blocks
+share it. Mappings retained in the loader cache continue to count.
 
 Components are considered in `diffusion`, `te`, `vae` order so that repeatedly
 used diffusion weights have priority. Each component's weights use the first

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_extend_backend.cpp

@@ -13,6 +13,7 @@
 #endif
 
 #include "core/util.h"
+#include "ggml-backend-impl.h"
 #include "ggml-impl.h"
 #include "stable-diffusion.h"
 
@@ -433,6 +434,24 @@ bool sd_backend_is_cpu(ggml_backend_t backend) {
     return dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU;
 }
 
+ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device,
+                                                          void* ptr,
+                                                          size_t size,
+                                                          size_t max_tensor_size) {
+    ggml_backend_buffer_t buffer = ggml_backend_dev_buffer_from_host_ptr(device, ptr, size, max_tensor_size);
+    if (buffer != nullptr && buffer->context == nullptr) {
+        ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(device);
+        if (reg != nullptr && std::strcmp(ggml_backend_reg_name(reg), "Metal") == 0) {
+            // Metal can wrap a failed mapping in a non-null buffer. Its free callback also
+            // dereferences the missing context, so only release the outer buffer.
+            buffer->iface.free_buffer = nullptr;
+            ggml_backend_buffer_free(buffer);
+            return nullptr;
+        }
+    }
+    return buffer;
+}
+
 bool sd_backend_supports_cuda_mma(ggml_backend_t backend) {
 #ifdef SD_USE_CUDA
     if (!sd_backend_is(backend, "CUDA")) {

src/core/ggml_extend_backend.h

@@ -88,6 +88,10 @@ private:
 bool sd_backend_is(ggml_backend_t backend, const std::string& name);
 bool sd_backend_is_cpu(ggml_backend_t backend);
 bool sd_backend_supports_cuda_mma(ggml_backend_t backend);
+ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device,
+                                                          void* ptr,
+                                                          size_t size,
+                                                          size_t max_tensor_size);
 ggml_backend_t sd_backend_cpu_init();
 bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads);
 ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,

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/common/rope.hpp

@@ -5,6 +5,7 @@
 #include <cassert>
 #include <cmath>
 #include <set>
+#include <utility>
 #include <vector>
 #include "core/ggml_extend.h"
 #include "core/ggml_runner.h"
@@ -16,6 +17,45 @@ namespace Rope {
         ErnieImage,
     };
 
+    struct SpatialRegion {
+        size_t begin;
+        size_t count;
+        float height_period;
+        float width_period;
+        int height_axis = 1;
+        int width_axis  = 2;
+    };
+
+    struct PositionLayout {
+        // Token ranges are relative to one batch item.
+        std::vector<SpatialRegion> images;
+        size_t token_count = 0;
+
+        void append_tokens(size_t count) {
+            token_count += count;
+        }
+
+        void append_image(int height, int width, int frames = 1, float height_step = 1.f, float width_step = 1.f) {
+            size_t count = static_cast<size_t>(height) * width * frames;
+            images.push_back({token_count, count, height * height_step, width * width_step});
+            append_tokens(count);
+        }
+    };
+
+    struct Frequency {
+        size_t axis;
+        float omega;
+    };
+
+    struct Embedding {
+        std::vector<float> values;
+        std::vector<std::vector<float>> ids;
+        PositionLayout positions;
+        std::vector<Frequency> frequencies;
+        EmbedNDLayout layout = EmbedNDLayout::Matrix;
+        int batch_size       = 1;
+    };
+
     enum class RefIndexMode {
         FIXED,
         INCREASE,
@@ -56,40 +96,25 @@ namespace Rope {
         return flat_vec;
     }
 
-    __STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
-                                                           int dim,
-                                                           float theta,
-                                                           const std::vector<int>& axis_wrap_dims = {}) {
+    __STATIC_INLINE__ std::vector<float> rope_frequencies(int dim, float theta) {
         assert(dim % 2 == 0);
-        int half_dim = dim / 2;
-
+        int half_dim             = dim / 2;
         std::vector<float> scale = linspace(0.f, (dim * 1.f - 2) / dim, half_dim);
-
         std::vector<float> omega(half_dim);
         for (int i = 0; i < half_dim; ++i) {
             omega[i] = 1.0f / ::powf(1.f * theta, scale[i]);
         }
+        return omega;
+    }
 
+    __STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
+                                                           const std::vector<float>& omega) {
+        int half_dim    = static_cast<int>(omega.size());
         size_t pos_size = pos.size();
         std::vector<std::vector<float>> out(pos_size, std::vector<float>(half_dim));
         for (size_t i = 0; i < pos_size; ++i) {
             for (size_t j = 0; j < half_dim; ++j) {
                 float angle = pos[i] * omega[j];
-                if (!axis_wrap_dims.empty()) {
-                    size_t wrap_size = axis_wrap_dims.size();
-                    // mod batch size since we only store this for one item in the batch
-                    size_t wrap_idx = wrap_size > 0 ? (i % wrap_size) : 0;
-                    int wrap_dim    = axis_wrap_dims[wrap_idx];
-                    if (wrap_dim > 0) {
-                        constexpr float TWO_PI = 6.28318530717958647692f;
-                        float cycles           = omega[j] * wrap_dim / TWO_PI;
-                        // closest periodic harmonic, necessary to ensure things neatly tile
-                        // without this round, things don't tile at the boundaries and you end up
-                        // with the model knowing what is "center"
-                        float rounded = std::round(cycles);
-                        angle         = pos[i] * TWO_PI * rounded / wrap_dim;
-                    }
-                }
 
                 out[i][j] = angle;
             }
@@ -108,6 +133,12 @@ namespace Rope {
         return result;
     }
 
+    __STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
+                                                           int dim,
+                                                           float theta) {
+        return rope(pos, rope_frequencies(dim, theta));
+    }
+
     // Generate IDs for image patches and text
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_flux_txt_ids(int bs, int context_len, int axes_dim_num, std::set<int> arange_dims) {
         auto txt_ids = std::vector<std::vector<float>>(bs * context_len, std::vector<float>(axes_dim_num, 0.0f));
@@ -136,12 +167,16 @@ namespace Rope {
                                                                        int patch_size,
                                                                        int bs,
                                                                        int axes_dim_num,
-                                                                       int index       = 0,
-                                                                       int h_offset    = 0,
-                                                                       int w_offset    = 0,
-                                                                       bool scale_rope = false) {
+                                                                       int index              = 0,
+                                                                       int h_offset           = 0,
+                                                                       int w_offset           = 0,
+                                                                       bool scale_rope        = false,
+                                                                       PositionLayout* layout = nullptr) {
         int h_len = (h + (patch_size / 2)) / patch_size;
         int w_len = (w + (patch_size / 2)) / patch_size;
+        if (layout) {
+            layout->append_image(h_len, w_len);
+        }
         std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(axes_dim_num, 0.0));
 
         int h_start = h_offset;
@@ -192,8 +227,8 @@ namespace Rope {
                                                   int bs,
                                                   const std::vector<float>& axis_thetas,
                                                   const std::vector<int>& axes_dim,
-                                                  const std::vector<std::vector<int>>& wrap_dims = {},
-                                                  EmbedNDLayout layout                           = EmbedNDLayout::Matrix) {
+                                                  EmbedNDLayout layout                = EmbedNDLayout::Matrix,
+                                                  std::vector<Frequency>* frequencies = nullptr) {
         std::vector<std::vector<float>> trans_ids = transpose(ids);
         size_t pos_len                            = ids.size() / bs;
         size_t num_axes                           = axes_dim.size();
@@ -205,19 +240,25 @@ namespace Rope {
         for (int d : axes_dim)
             emb_dim += d / 2;
 
+        if (frequencies) {
+            frequencies->clear();
+            frequencies->reserve(emb_dim);
+        }
         std::vector<std::vector<float>> emb(bs * pos_len, std::vector<float>(emb_dim * 2 * 2, 0.0));
         size_t offset = 0;
         for (size_t i = 0; i < num_axes; ++i) {
-            std::vector<int> axis_wrap_dims;
-            if (!wrap_dims.empty() && i < (int)wrap_dims.size()) {
-                axis_wrap_dims = wrap_dims[i];
-            }
             float axis_theta = 10000.0f;
             if (!axis_thetas.empty()) {
                 axis_theta = axis_thetas[std::min(i, axis_thetas.size() - 1)];
             }
+            auto omega = rope_frequencies(axes_dim[i], axis_theta);
+            if (frequencies) {
+                for (float frequency : omega) {
+                    frequencies->push_back({i, frequency});
+                }
+            }
             std::vector<std::vector<float>> rope_emb =
-                rope(trans_ids[i], axes_dim[i], axis_theta, axis_wrap_dims);  // [bs*pos_len, axes_dim[i]/2 * 2 * 2]
+                rope(trans_ids[i], omega);  // [bs*pos_len, axes_dim[i]/2 * 2 * 2]
             for (int b = 0; b < bs; ++b) {
                 for (int j = 0; j < pos_len; ++j) {
                     for (int k = 0; k < rope_emb[0].size(); ++k) {
@@ -253,10 +294,10 @@ namespace Rope {
                                                   int bs,
                                                   float theta,
                                                   const std::vector<int>& axes_dim,
-                                                  const std::vector<std::vector<int>>& wrap_dims = {},
-                                                  EmbedNDLayout layout                           = EmbedNDLayout::Matrix) {
+                                                  EmbedNDLayout layout                = EmbedNDLayout::Matrix,
+                                                  std::vector<Frequency>* frequencies = nullptr) {
         std::vector<float> axis_thetas(axes_dim.size(), theta);
-        return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims, layout);
+        return embed_nd(ids, bs, axis_thetas, axes_dim, layout, frequencies);
     }
 
     __STATIC_INLINE__ std::vector<float> embed_interleaved_mrope(const std::vector<std::vector<float>>& ids,
@@ -264,7 +305,7 @@ namespace Rope {
                                                                  float theta,
                                                                  int head_dim,
                                                                  const std::vector<int>& mrope_section,
-                                                                 const std::vector<std::vector<int>>& axis_wrap_dims = {}) {
+                                                                 std::vector<Frequency>* frequencies = nullptr) {
         GGML_ASSERT(bs > 0);
         GGML_ASSERT(head_dim % 2 == 0);
         GGML_ASSERT(mrope_section.size() >= 3);
@@ -273,20 +314,26 @@ namespace Rope {
         size_t pos_len                            = ids.size() / bs;
         int half_dim                              = head_dim / 2;
 
+        auto omega = rope_frequencies(head_dim, theta);
+        if (frequencies) {
+            frequencies->clear();
+            for (float frequency : omega) {
+                frequencies->push_back({0, frequency});
+            }
+        }
         std::vector<std::vector<std::vector<float>>> axis_embs;
         axis_embs.reserve(3);
         for (int axis = 0; axis < 3; ++axis) {
-            std::vector<int> axis_wrap;
-            if (axis < static_cast<int>(axis_wrap_dims.size())) {
-                axis_wrap = axis_wrap_dims[axis];
-            }
-            axis_embs.push_back(rope(trans_ids[axis], head_dim, theta, axis_wrap));
+            axis_embs.push_back(rope(trans_ids[axis], omega));
         }
 
         std::vector<std::vector<float>> emb = axis_embs[0];
         for (int axis = 1; axis < 3; ++axis) {
             int length = std::min<int>(mrope_section[axis] * 3, half_dim);
             for (int freq_idx = axis; freq_idx < length; freq_idx += 3) {
+                if (frequencies) {
+                    (*frequencies)[freq_idx].axis = axis;
+                }
                 for (size_t pos_idx = 0; pos_idx < bs * pos_len; ++pos_idx) {
                     for (int k = 0; k < 4; ++k) {
                         emb[pos_idx][4 * freq_idx + k] = axis_embs[axis][pos_idx][4 * freq_idx + k];
@@ -298,13 +345,13 @@ namespace Rope {
         return flatten(emb);
     }
 
-    __STATIC_INLINE__ std::vector<float> embed_2d_interleaved(int height,
-                                                              int width,
-                                                              int dim,
-                                                              float theta    = 10000.f,
-                                                              float scale    = 16.f,
-                                                              int ref_grid_h = 0,
-                                                              int ref_grid_w = 0) {
+    __STATIC_INLINE__ Embedding embed_2d_interleaved(int height,
+                                                     int width,
+                                                     int dim,
+                                                     float theta    = 10000.f,
+                                                     float scale    = 16.f,
+                                                     int ref_grid_h = 0,
+                                                     int ref_grid_w = 0) {
         assert(dim % 4 == 0);
         int half_dim      = dim / 2;
         int dim_axis      = dim / 2;
@@ -318,6 +365,10 @@ namespace Rope {
             w_ntk       = std::pow(static_cast<float>(width) / static_cast<float>(ref_grid_w), power);
         }
 
+        Embedding result;
+        result.positions.append_image(height, width, 1,
+                                      height > 1 ? scale / (height - 1) : 1.f,
+                                      width > 1 ? scale / (width - 1) : 1.f);
         std::vector<float> x_pos;
         std::vector<float> y_pos;
         x_pos.reserve(static_cast<size_t>(height) * width);
@@ -326,13 +377,20 @@ namespace Rope {
             float y = height == 1 ? 0.f : scale * static_cast<float>(iy) / static_cast<float>(height - 1);
             for (int ix = 0; ix < width; ++ix) {
                 float x = width == 1 ? 0.f : scale * static_cast<float>(ix) / static_cast<float>(width - 1);
+                result.ids.push_back({0.f, y, x});
                 x_pos.push_back(x);
                 y_pos.push_back(y);
             }
         }
 
-        auto x_emb = rope(x_pos, dim_axis, theta * w_ntk);
-        auto y_emb = rope(y_pos, dim_axis, theta * h_ntk);
+        auto x_freq = rope_frequencies(dim_axis, theta * w_ntk);
+        auto y_freq = rope_frequencies(dim_axis, theta * h_ntk);
+        auto x_emb  = rope(x_pos, x_freq);
+        auto y_emb  = rope(y_pos, y_freq);
+        for (int i = 0; i < axis_half_dim; ++i) {
+            result.frequencies.push_back({2, x_freq[i]});
+            result.frequencies.push_back({1, y_freq[i]});
+        }
 
         std::vector<float> out(static_cast<size_t>(height) * width * half_dim * 4);
         for (int pos = 0; pos < height * width; ++pos) {
@@ -348,7 +406,8 @@ namespace Rope {
                 }
             }
         }
-        return out;
+        result.values = std::move(out);
+        return result;
     }
 
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_refs_ids(int patch_size,
@@ -359,7 +418,8 @@ namespace Rope {
                                                                    RefIndexMode ref_index_mode,
                                                                    float ref_index_scale,
                                                                    bool scale_rope,
-                                                                   int base_offset = 0) {
+                                                                   int base_offset        = 0,
+                                                                   PositionLayout* layout = nullptr) {
         std::vector<std::vector<float>> ids;
         int curr_h_offset = 0;
         int curr_w_offset = 0;
@@ -386,7 +446,8 @@ namespace Rope {
                                             static_cast<int>(index * ref_index_scale),
                                             h_offset + base_offset,
                                             w_offset + base_offset,
-                                            scale_rope);
+                                            scale_rope,
+                                            layout);
             ids          = concat_ids(ids, ref_ids, bs);
 
             if (ref_index_mode == RefIndexMode::INCREASE) {
@@ -409,88 +470,53 @@ namespace Rope {
                                                                    const std::vector<ggml_tensor*>& ref_latents,
                                                                    RefIndexMode ref_index_mode,
                                                                    float ref_index_scale,
-                                                                   bool is_longcat) {
+                                                                   bool is_longcat,
+                                                                   PositionLayout* layout = nullptr) {
+        if (layout) {
+            layout->append_tokens(context_len);
+        }
         int x_index = is_longcat ? 1 : 0;
 
         auto txt_ids = is_longcat ? gen_longcat_txt_ids(bs, context_len, axes_dim_num) : gen_flux_txt_ids(bs, context_len, axes_dim_num, txt_arange_dims);
         int offset   = is_longcat ? context_len : 0;
-        auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, x_index, offset, offset);
+        auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, x_index, offset, offset, false, layout);
 
         auto ids = concat_ids(txt_ids, img_ids, bs);
         if (ref_latents.size() > 0) {
-            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, ref_index_mode, ref_index_scale, false, offset);
+            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, ref_index_mode, ref_index_scale, false, offset, layout);
             ids           = concat_ids(ids, refs_ids, bs);
         }
         return ids;
     }
 
     // Generate flux positional embeddings
-    __STATIC_INLINE__ std::vector<float> gen_flux_pe(int h,
-                                                     int w,
-                                                     int patch_size,
-                                                     int bs,
-                                                     int context_len,
-                                                     std::set<int> txt_arange_dims,
-                                                     const std::vector<ggml_tensor*>& ref_latents,
-                                                     RefIndexMode ref_index_mode,
-                                                     float ref_index_scale,
-                                                     int theta,
-                                                     bool circular_h,
-                                                     bool circular_w,
-                                                     const std::vector<int>& axes_dim,
-                                                     bool is_longcat) {
-        std::vector<std::vector<float>> ids = gen_flux_ids(h,
-                                                           w,
-                                                           patch_size,
-                                                           bs,
-                                                           static_cast<int>(axes_dim.size()),
-                                                           context_len,
-                                                           txt_arange_dims,
-                                                           ref_latents,
-                                                           ref_index_mode,
-                                                           ref_index_scale,
-                                                           is_longcat);
-        std::vector<std::vector<int>> wrap_dims;
-        if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
-            int h_len = (h + (patch_size / 2)) / patch_size;
-            int w_len = (w + (patch_size / 2)) / patch_size;
-            if (h_len > 0 && w_len > 0) {
-                size_t pos_len = ids.size() / bs;
-                wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
-                size_t cursor           = context_len;  // text first
-                const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
-                for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
-                    if (circular_h) {
-                        wrap_dims[1][cursor + token_i] = h_len;
-                    }
-                    if (circular_w) {
-                        wrap_dims[2][cursor + token_i] = w_len;
-                    }
-                }
-                cursor += img_tokens;
-                // reference latents
-                for (ggml_tensor* ref : ref_latents) {
-                    if (ref == nullptr) {
-                        continue;
-                    }
-                    int ref_h         = static_cast<int>(ref->ne[1]);
-                    int ref_w         = static_cast<int>(ref->ne[0]);
-                    int ref_h_l       = (ref_h + (patch_size / 2)) / patch_size;
-                    int ref_w_l       = (ref_w + (patch_size / 2)) / patch_size;
-                    size_t ref_tokens = static_cast<size_t>(ref_h_l) * static_cast<size_t>(ref_w_l);
-                    for (size_t token_i = 0; token_i < ref_tokens; ++token_i) {
-                        if (circular_h) {
-                            wrap_dims[1][cursor + token_i] = ref_h_l;
-                        }
-                        if (circular_w) {
-                            wrap_dims[2][cursor + token_i] = ref_w_l;
-                        }
-                    }
-                    cursor += ref_tokens;
-                }
-            }
-        }
-        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
+    __STATIC_INLINE__ Embedding gen_flux_pe(int h,
+                                            int w,
+                                            int patch_size,
+                                            int bs,
+                                            int context_len,
+                                            std::set<int> txt_arange_dims,
+                                            const std::vector<ggml_tensor*>& ref_latents,
+                                            RefIndexMode ref_index_mode,
+                                            float ref_index_scale,
+                                            int theta,
+                                            const std::vector<int>& axes_dim,
+                                            bool is_longcat) {
+        Embedding result;
+        result.batch_size = bs;
+        result.ids        = gen_flux_ids(h,
+                                         w,
+                                         patch_size,
+                                         bs,
+                                         static_cast<int>(axes_dim.size()),
+                                         context_len,
+                                         txt_arange_dims,
+                                         ref_latents,
+                                         ref_index_mode,
+                                         ref_index_scale,
+                                         is_longcat, &result.positions);
+        result.values     = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
+        return result;
     }
 
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
@@ -500,14 +526,18 @@ namespace Rope {
                                                                   int ph,
                                                                   int pw,
                                                                   int bs,
-                                                                  int t_offset    = 0,
-                                                                  int h_offset    = 0,
-                                                                  int w_offset    = 0,
-                                                                  bool scale_rope = false) {
+                                                                  int t_offset           = 0,
+                                                                  int h_offset           = 0,
+                                                                  int w_offset           = 0,
+                                                                  bool scale_rope        = false,
+                                                                  PositionLayout* layout = nullptr) {
         int t_len = (t + (pt / 2)) / pt;
         int h_len = (h + (ph / 2)) / ph;
         int w_len = (w + (pw / 2)) / pw;
 
+        if (layout) {
+            layout->append_image(h_len, w_len, t_len);
+        }
         std::vector<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
 
         if (scale_rope) {
@@ -573,7 +603,11 @@ namespace Rope {
                                                                          int bs,
                                                                          int context_len,
                                                                          const std::vector<ggml_tensor*>& ref_latents,
-                                                                         RefIndexMode ref_index_mode) {
+                                                                         RefIndexMode ref_index_mode,
+                                                                         PositionLayout* layout = nullptr) {
+        if (layout) {
+            layout->append_tokens(context_len);
+        }
         int h_len        = (h + (patch_size / 2)) / patch_size;
         int w_len        = (w + (patch_size / 2)) / patch_size;
         int txt_id_start = std::max(h_len, w_len) / 2;
@@ -585,90 +619,49 @@ namespace Rope {
             }
         }
         int axes_dim_num = 3;
-        auto img_ids     = gen_vid_ids(t, h, w, 1, patch_size, patch_size, bs, 0, 0, 0, true);
+        auto img_ids     = gen_vid_ids(t, h, w, 1, patch_size, patch_size, bs, 0, 0, 0, true, layout);
         auto ids         = concat_ids(txt_ids_repeated, img_ids, bs);
         if (ref_latents.size() > 0) {
             int ref_start_index = ref_index_mode == RefIndexMode::DECREASE ? 0 : 1;
-            auto refs_ids       = gen_refs_ids(patch_size, bs, axes_dim_num, ref_start_index, ref_latents, ref_index_mode, 1.f, true);
+            auto refs_ids       = gen_refs_ids(patch_size, bs, axes_dim_num, ref_start_index, ref_latents, ref_index_mode, 1.f, true, 0, layout);
             ids                 = concat_ids(ids, refs_ids, bs);
         }
         return ids;
     }
 
     // Generate qwen_image positional embeddings
-    __STATIC_INLINE__ std::vector<float> gen_qwen_image_pe(int t,
-                                                           int h,
-                                                           int w,
-                                                           int patch_size,
-                                                           int bs,
-                                                           int context_len,
-                                                           const std::vector<ggml_tensor*>& ref_latents,
-                                                           RefIndexMode ref_index_mode,
-                                                           int theta,
-                                                           bool circular_h,
-                                                           bool circular_w,
-                                                           const std::vector<int>& axes_dim) {
-        std::vector<std::vector<float>> ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode);
-        std::vector<std::vector<int>> wrap_dims;
-        // This logic simply stores the (pad and patch_adjusted) sizes of images so we can make sure rope correctly tiles
-        if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
-            int pad_h = (patch_size - (h % patch_size)) % patch_size;
-            int pad_w = (patch_size - (w % patch_size)) % patch_size;
-            int h_len = (h + pad_h) / patch_size;
-            int w_len = (w + pad_w) / patch_size;
-            if (h_len > 0 && w_len > 0) {
-                const size_t total_tokens = ids.size();
-                // Track per-token wrap lengths for the row/column axes so only spatial tokens become periodic.
-                wrap_dims.assign(axes_dim.size(), std::vector<int>(total_tokens / bs, 0));
-                size_t cursor           = context_len;  // ignore text tokens
-                const size_t img_tokens = static_cast<size_t>(t) * static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
-                for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
-                    if (circular_h) {
-                        wrap_dims[1][cursor + token_i] = h_len;
-                    }
-                    if (circular_w) {
-                        wrap_dims[2][cursor + token_i] = w_len;
-                    }
-                }
-                cursor += img_tokens;
-                // For each reference image, store wrap sizes as well
-                for (ggml_tensor* ref : ref_latents) {
-                    if (ref == nullptr) {
-                        continue;
-                    }
-                    int ref_h           = static_cast<int>(ref->ne[1]);
-                    int ref_w           = static_cast<int>(ref->ne[0]);
-                    int ref_pad_h       = (patch_size - (ref_h % patch_size)) % patch_size;
-                    int ref_pad_w       = (patch_size - (ref_w % patch_size)) % patch_size;
-                    int ref_h_len       = (ref_h + ref_pad_h) / patch_size;
-                    int ref_w_len       = (ref_w + ref_pad_w) / patch_size;
-                    size_t ref_n_tokens = static_cast<size_t>(ref_h_len) * static_cast<size_t>(ref_w_len);
-                    for (size_t token_i = 0; token_i < ref_n_tokens; ++token_i) {
-                        if (circular_h) {
-                            wrap_dims[1][cursor + token_i] = ref_h_len;
-                        }
-                        if (circular_w) {
-                            wrap_dims[2][cursor + token_i] = ref_w_len;
-                        }
-                    }
-                    cursor += ref_n_tokens;
-                }
-            }
-        }
-        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
+    __STATIC_INLINE__ Embedding gen_qwen_image_pe(int t,
+                                                  int h,
+                                                  int w,
+                                                  int patch_size,
+                                                  int bs,
+                                                  int context_len,
+                                                  const std::vector<ggml_tensor*>& ref_latents,
+                                                  RefIndexMode ref_index_mode,
+                                                  int theta,
+                                                  const std::vector<int>& axes_dim) {
+        Embedding result;
+        result.batch_size = bs;
+        result.ids        = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode, &result.positions);
+        result.values     = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
+        return result;
     }
 
-    __STATIC_INLINE__ std::vector<float> gen_mage_flow_pe(int h,
-                                                          int w,
-                                                          int bs,
-                                                          int context_len,
-                                                          const std::vector<ggml_tensor*>& ref_latents,
-                                                          int theta,
-                                                          const std::vector<int>& axes_dim) {
+    __STATIC_INLINE__ Embedding gen_mage_flow_pe(int h,
+                                                 int w,
+                                                 int bs,
+                                                 int context_len,
+                                                 const std::vector<ggml_tensor*>& ref_latents,
+                                                 int theta,
+                                                 const std::vector<int>& axes_dim) {
+        Embedding result;
+        result.batch_size = bs;
+        result.positions.append_tokens(context_len);
         const int axes_dim_num = static_cast<int>(axes_dim.size());
-        auto make_image_ids    = [=](int image_h, int image_w, int image_index) {
+        auto make_image_ids    = [=, &result](int image_h, int image_w, int image_index) {
             std::vector<std::vector<float>> image_ids(static_cast<size_t>(bs) * image_h * image_w,
                                                          std::vector<float>(axes_dim_num, 0.f));
+            result.positions.append_image(image_h, image_w);
             int h_start = -(image_h - image_h / 2);
             int w_start = -(image_w - image_w / 2);
             for (int b = 0; b < bs; ++b) {
@@ -692,15 +685,18 @@ namespace Rope {
                                           static_cast<int>(i + 1));
             ids          = concat_ids(ids, ref_ids, bs);
         }
-        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
+        result.ids    = std::move(ids);
+        result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
+        return result;
     }
 
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_lens_ids(int h,
                                                                    int w,
                                                                    int bs,
                                                                    int context_len,
-                                                                   bool scale_rope = true) {
-        auto img_ids_repeated = gen_flux_img_ids(h, w, 1, bs, 3, 0, 0, 0, scale_rope);
+                                                                   bool scale_rope        = true,
+                                                                   PositionLayout* layout = nullptr) {
+        auto img_ids_repeated = gen_flux_img_ids(h, w, 1, bs, 3, 0, 0, 0, scale_rope, layout);
 
         int txt_id_start = scale_rope ? std::max(h / 2, w / 2) : 0;
         auto txt_ids     = linspace<float>(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len);
@@ -711,44 +707,37 @@ namespace Rope {
             }
         }
 
+        if (layout) {
+            layout->append_tokens(context_len);
+        }
         return concat_ids(img_ids_repeated, txt_ids_repeated, bs);
     }
 
-    __STATIC_INLINE__ std::vector<float> gen_lens_pe(int h,
-                                                     int w,
-                                                     int bs,
-                                                     int context_len,
-                                                     int theta,
-                                                     bool circular_h,
-                                                     bool circular_w,
-                                                     const std::vector<int>& axes_dim) {
-        std::vector<std::vector<float>> ids = gen_lens_ids(h, w, bs, context_len, true);
-        std::vector<std::vector<int>> wrap_dims;
-        if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
-            size_t pos_len = ids.size() / bs;
-            wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
-            const size_t img_tokens = static_cast<size_t>(h) * static_cast<size_t>(w);
-            for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
-                if (circular_h) {
-                    wrap_dims[1][token_i] = h;
-                }
-                if (circular_w) {
-                    wrap_dims[2][token_i] = w;
-                }
-            }
-        }
-
-        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
+    __STATIC_INLINE__ Embedding gen_lens_pe(int h,
+                                            int w,
+                                            int bs,
+                                            int context_len,
+                                            int theta,
+                                            const std::vector<int>& axes_dim) {
+        Embedding result;
+        result.batch_size = bs;
+        result.ids        = gen_lens_ids(h, w, bs, context_len, true, &result.positions);
+        result.values     = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
+        return result;
     }
 
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_ernie_image_ids(int h,
                                                                           int w,
                                                                           int patch_size,
                                                                           int bs,
-                                                                          int context_len) {
+                                                                          int context_len,
+                                                                          PositionLayout* layout = nullptr) {
         int h_len = h / patch_size;
         int w_len = w / patch_size;
 
+        if (layout) {
+            layout->append_image(h_len, w_len);
+        }
         std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(3, 0.0f));
         std::vector<float> h_ids = linspace<float>(0.f, static_cast<float>(h_len - 1), h_len);
         std::vector<float> w_ids = linspace<float>(0.f, static_cast<float>(w_len - 1), w_len);
@@ -774,39 +763,25 @@ namespace Rope {
             }
         }
 
+        if (layout) {
+            layout->append_tokens(context_len);
+        }
         return concat_ids(img_ids_repeated, txt_ids, bs);
     }
 
-    __STATIC_INLINE__ std::vector<float> gen_ernie_image_pe(int h,
-                                                            int w,
-                                                            int patch_size,
-                                                            int bs,
-                                                            int context_len,
-                                                            int theta,
-                                                            bool circular_h,
-                                                            bool circular_w,
-                                                            const std::vector<int>& axes_dim) {
-        std::vector<std::vector<float>> ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len);
-        std::vector<std::vector<int>> wrap_dims;
-        if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
-            int h_len = h / patch_size;
-            int w_len = w / patch_size;
-            if (h_len > 0 && w_len > 0) {
-                size_t pos_len = ids.size() / bs;
-                wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
-                const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
-                for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
-                    if (circular_h) {
-                        wrap_dims[1][token_i] = h_len;
-                    }
-                    if (circular_w) {
-                        wrap_dims[2][token_i] = w_len;
-                    }
-                }
-            }
-        }
-
-        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage);
+    __STATIC_INLINE__ Embedding gen_ernie_image_pe(int h,
+                                                   int w,
+                                                   int patch_size,
+                                                   int bs,
+                                                   int context_len,
+                                                   int theta,
+                                                   const std::vector<int>& axes_dim) {
+        Embedding result;
+        result.batch_size = bs;
+        result.layout     = EmbedNDLayout::ErnieImage;
+        result.ids        = gen_ernie_image_ids(h, w, patch_size, bs, context_len, &result.positions);
+        result.values     = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
+        return result;
     }
 
     // Generate wan positional embeddings
@@ -905,7 +880,8 @@ namespace Rope {
                                                                       int context_len,
                                                                       int seq_multi_of,
                                                                       const std::vector<ggml_tensor*>& ref_latents,
-                                                                      RefIndexMode ref_index_mode) {
+                                                                      RefIndexMode ref_index_mode,
+                                                                      PositionLayout* layout = nullptr) {
         SD_UNUSED(ref_index_mode);
         int padded_context_len = context_len + bound_mod(context_len, seq_multi_of);
         auto txt_ids           = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
@@ -913,11 +889,17 @@ namespace Rope {
             txt_ids[i][0] = (i % padded_context_len) + 1.f;
         }
 
+        if (layout) {
+            layout->append_tokens(padded_context_len);
+        }
         int axes_dim_num = 3;
         int index        = padded_context_len + 1;
-        auto img_ids     = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index);
+        auto img_ids     = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index, 0, 0, false, layout);
 
         int img_pad_len = bound_mod(static_cast<int>(img_ids.size() / bs), seq_multi_of);
+        if (layout) {
+            layout->append_tokens(img_pad_len);
+        }
         if (img_pad_len > 0) {
             std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f));
             img_ids = concat_ids(img_ids, img_pad_ids, bs);
@@ -936,7 +918,8 @@ namespace Rope {
                                                                           int patch_size,
                                                                           int bs,
                                                                           int context_len,
-                                                                          int seq_multi_of) {
+                                                                          int seq_multi_of,
+                                                                          PositionLayout* layout = nullptr) {
         int context_pad_len    = bound_mod(context_len, seq_multi_of);
         int padded_context_len = context_len + context_pad_len;
         auto txt_ids           = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
@@ -947,11 +930,17 @@ namespace Rope {
             }
         }
 
+        if (layout) {
+            layout->append_tokens(padded_context_len);
+        }
         int axes_dim_num = 3;
         int index        = padded_context_len + 1;
-        auto img_ids     = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index);
+        auto img_ids     = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index, 0, 0, false, layout);
 
         int img_pad_len = bound_mod(static_cast<int>(img_ids.size() / bs), seq_multi_of);
+        if (layout) {
+            layout->append_tokens(img_pad_len);
+        }
         if (img_pad_len > 0) {
             std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f));
             img_ids = concat_ids(img_ids, img_pad_ids, bs);
@@ -968,7 +957,8 @@ namespace Rope {
                                                                                int patch_size,
                                                                                int context_len,
                                                                                int sigvq_len,
-                                                                               int seq_multi_of) {
+                                                                               int seq_multi_of,
+                                                                               PositionLayout* layout = nullptr) {
         const int context_pad    = bound_mod(context_len, seq_multi_of);
         const int padded_context = context_len + context_pad;
         const int h_len          = (h + (patch_size / 2)) / patch_size;
@@ -994,11 +984,17 @@ namespace Rope {
             cursor += 2;
         }
 
+        if (layout) {
+            layout->append_tokens(cap_ids.size());
+        }
         std::vector<std::vector<float>> img_ids;
         for (int copy = 0; copy < 2; ++copy) {
-            auto ids = gen_flux_img_ids(h, w, patch_size, 1, 3, cap_end_positions[copy]);
+            auto ids = gen_flux_img_ids(h, w, patch_size, 1, 3, cap_end_positions[copy], 0, 0, false, layout);
             img_ids.insert(img_ids.end(), ids.begin(), ids.end());
             img_ids.insert(img_ids.end(), image_pad, std::vector<float>(3, 0.f));
+            if (layout) {
+                layout->append_tokens(image_pad);
+            }
         }
 
         const int sigvq_start = static_cast<int>(cap_ids.size() + img_ids.size()) + 1;
@@ -1016,95 +1012,59 @@ namespace Rope {
         ids.insert(ids.end(), cap_ids.begin(), cap_ids.end());
         ids.insert(ids.end(), img_ids.begin(), img_ids.end());
         ids.insert(ids.end(), sigvq_ids.begin(), sigvq_ids.end());
+        if (layout) {
+            layout->append_tokens(sigvq_ids.size());
+        }
         SD_UNUSED(padded_image);
         return ids;
     }
 
-    __STATIC_INLINE__ std::vector<float> gen_llada_image_edit_pe(int h,
-                                                                 int w,
-                                                                 int patch_size,
-                                                                 int context_len,
-                                                                 int sigvq_len,
-                                                                 int seq_multi_of,
-                                                                 int theta,
-                                                                 const std::vector<int>& axes_dim) {
-        auto ids = gen_llada_image_edit_ids(h, w, patch_size, context_len, sigvq_len, seq_multi_of);
-        return embed_nd(ids, 1, static_cast<float>(theta), axes_dim, {});
+    __STATIC_INLINE__ Embedding gen_llada_image_edit_pe(int h,
+                                                        int w,
+                                                        int patch_size,
+                                                        int context_len,
+                                                        int sigvq_len,
+                                                        int seq_multi_of,
+                                                        int theta,
+                                                        const std::vector<int>& axes_dim) {
+        Embedding result;
+        result.batch_size = 1;
+        result.ids        = gen_llada_image_edit_ids(h, w, patch_size, context_len, sigvq_len, seq_multi_of, &result.positions);
+        result.values     = embed_nd(result.ids, 1, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
+        return result;
     }
 
-    __STATIC_INLINE__ std::vector<float> gen_llada_image_pe(int h,
-                                                            int w,
-                                                            int patch_size,
-                                                            int bs,
-                                                            int context_len,
-                                                            int seq_multi_of,
-                                                            int theta,
-                                                            bool circular_h,
-                                                            bool circular_w,
-                                                            const std::vector<int>& axes_dim) {
-        std::vector<std::vector<float>> ids = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of);
-        std::vector<std::vector<int>> wrap_dims;
-        if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
-            int pad_h = (patch_size - (h % patch_size)) % patch_size;
-            int pad_w = (patch_size - (w % patch_size)) % patch_size;
-            int h_len = (h + pad_h) / patch_size;
-            int w_len = (w + pad_w) / patch_size;
-            if (h_len > 0 && w_len > 0) {
-                size_t pos_len = ids.size() / bs;
-                wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
-                size_t cursor     = context_len + bound_mod(context_len, seq_multi_of);
-                size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
-                for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
-                    if (circular_h) {
-                        wrap_dims[1][cursor + token_i] = h_len;
-                    }
-                    if (circular_w) {
-                        wrap_dims[2][cursor + token_i] = w_len;
-                    }
-                }
-            }
-        }
-
-        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
+    __STATIC_INLINE__ Embedding gen_llada_image_pe(int h,
+                                                   int w,
+                                                   int patch_size,
+                                                   int bs,
+                                                   int context_len,
+                                                   int seq_multi_of,
+                                                   int theta,
+                                                   const std::vector<int>& axes_dim) {
+        Embedding result;
+        result.batch_size = bs;
+        result.ids        = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, &result.positions);
+        result.values     = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
+        return result;
     }
 
     // Generate z_image positional embeddings
-    __STATIC_INLINE__ std::vector<float> gen_z_image_pe(int h,
-                                                        int w,
-                                                        int patch_size,
-                                                        int bs,
-                                                        int context_len,
-                                                        int seq_multi_of,
-                                                        const std::vector<ggml_tensor*>& ref_latents,
-                                                        RefIndexMode ref_index_mode,
-                                                        int theta,
-                                                        bool circular_h,
-                                                        bool circular_w,
-                                                        const std::vector<int>& axes_dim) {
-        std::vector<std::vector<float>> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode);
-        std::vector<std::vector<int>> wrap_dims;
-        if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
-            int pad_h = (patch_size - (h % patch_size)) % patch_size;
-            int pad_w = (patch_size - (w % patch_size)) % patch_size;
-            int h_len = (h + pad_h) / patch_size;
-            int w_len = (w + pad_w) / patch_size;
-            if (h_len > 0 && w_len > 0) {
-                size_t pos_len = ids.size() / bs;
-                wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
-                size_t cursor     = context_len + bound_mod(context_len, seq_multi_of);  // skip text (and its padding)
-                size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
-                for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
-                    if (circular_h) {
-                        wrap_dims[1][cursor + token_i] = h_len;
-                    }
-                    if (circular_w) {
-                        wrap_dims[2][cursor + token_i] = w_len;
-                    }
-                }
-            }
-        }
-
-        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
+    __STATIC_INLINE__ Embedding gen_z_image_pe(int h,
+                                               int w,
+                                               int patch_size,
+                                               int bs,
+                                               int context_len,
+                                               int seq_multi_of,
+                                               const std::vector<ggml_tensor*>& ref_latents,
+                                               RefIndexMode ref_index_mode,
+                                               int theta,
+                                               const std::vector<int>& axes_dim) {
+        Embedding result;
+        result.batch_size = bs;
+        result.ids        = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode, &result.positions);
+        result.values     = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
+        return result;
     }
 
     __STATIC_INLINE__ ggml_tensor* apply_rope(ggml_context* ctx,

src/model/common/rope_circular.hpp

@@ -0,0 +1,65 @@
+#ifndef __SD_MODEL_COMMON_ROPE_CIRCULAR_HPP__
+#define __SD_MODEL_COMMON_ROPE_CIRCULAR_HPP__
+
+#include "model/common/rope.hpp"
+
+namespace Rope {
+    __STATIC_INLINE__ void apply_circular(Embedding& embedding, bool circular_x, bool circular_y) {
+        if (!circular_x && !circular_y) {
+            return;
+        }
+
+        GGML_ASSERT(embedding.batch_size > 0);
+        GGML_ASSERT(embedding.ids.size() % embedding.batch_size == 0);
+        size_t pos_len  = embedding.ids.size() / embedding.batch_size;
+        size_t half_dim = embedding.frequencies.size();
+        GGML_ASSERT(embedding.positions.token_count == pos_len);
+        GGML_ASSERT(embedding.values.size() == embedding.ids.size() * half_dim * 4);
+
+        constexpr float TWO_PI = 6.28318530717958647692f;
+        for (const auto& region : embedding.positions.images) {
+            GGML_ASSERT(region.begin <= pos_len && region.count <= pos_len - region.begin);
+            for (size_t j = 0; j < half_dim; ++j) {
+                const auto& frequency = embedding.frequencies[j];
+                float period          = 0.f;
+                if (circular_y && frequency.axis == static_cast<size_t>(region.height_axis)) {
+                    period = region.height_period;
+                } else if (circular_x && frequency.axis == static_cast<size_t>(region.width_axis)) {
+                    period = region.width_period;
+                }
+                if (period <= 0) {
+                    continue;
+                }
+
+                // Quantize to periodic harmonics while preserving the original coordinate offsets.
+                float rounded = std::round(frequency.omega * period / TWO_PI);
+                for (int b = 0; b < embedding.batch_size; ++b) {
+                    size_t begin = b * pos_len + region.begin;
+                    for (size_t i = begin; i < begin + region.count; ++i) {
+                        GGML_ASSERT(frequency.axis < embedding.ids[i].size());
+                        float angle   = embedding.ids[i][frequency.axis] * TWO_PI * rounded / period;
+                        float cos_val = std::cos(angle);
+                        float sin_val = std::sin(angle);
+                        if (embedding.layout == EmbedNDLayout::ErnieImage) {
+                            size_t cos_offset                = (i * half_dim + j) * 2;
+                            size_t sin_offset                = embedding.ids.size() * half_dim * 2 + cos_offset;
+                            embedding.values[cos_offset]     = cos_val;
+                            embedding.values[cos_offset + 1] = cos_val;
+                            embedding.values[sin_offset]     = sin_val;
+                            embedding.values[sin_offset + 1] = sin_val;
+                        } else {
+                            size_t offset                = (i * half_dim + j) * 4;
+                            embedding.values[offset]     = cos_val;
+                            embedding.values[offset + 1] = -sin_val;
+                            embedding.values[offset + 2] = sin_val;
+                            embedding.values[offset + 3] = cos_val;
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+}  // namespace Rope
+
+#endif  // __SD_MODEL_COMMON_ROPE_CIRCULAR_HPP__

src/model/diffusion/anima.hpp

@@ -603,34 +603,37 @@ namespace Anima {
             return std::pow(extrapolation_ratio, static_cast<float>(axis_dim) / static_cast<float>(axis_dim - 2));
         }
 
-        static std::vector<float> gen_anima_image_pe_vec(int bs,
-                                                         int h,
-                                                         int w,
-                                                         int patch_size,
-                                                         int theta,
-                                                         const std::vector<int>& axes_dim,
-                                                         float h_extrapolation_ratio,
-                                                         float w_extrapolation_ratio,
-                                                         float t_extrapolation_ratio,
-                                                         const std::vector<ggml_tensor*>& ref_latents) {
-            auto ids = Rope::gen_flux_ids(h,
-                                          w,
-                                          patch_size,
-                                          bs,
-                                          static_cast<int>(axes_dim.size()),
-                                          0,
-                                          {},
-                                          ref_latents,
-                                          Rope::RefIndexMode::FIXED,
-                                          1.0f,
-                                          false);
+        static Rope::Embedding gen_anima_image_pe_vec(int bs,
+                                                      int h,
+                                                      int w,
+                                                      int patch_size,
+                                                      int theta,
+                                                      const std::vector<int>& axes_dim,
+                                                      float h_extrapolation_ratio,
+                                                      float w_extrapolation_ratio,
+                                                      float t_extrapolation_ratio,
+                                                      const std::vector<ggml_tensor*>& ref_latents) {
+            Rope::Embedding result;
+            result.batch_size = bs;
+            result.ids        = Rope::gen_flux_ids(h,
+                                                   w,
+                                                   patch_size,
+                                                   bs,
+                                                   static_cast<int>(axes_dim.size()),
+                                                   0,
+                                                   {},
+                                                   ref_latents,
+                                                   Rope::RefIndexMode::FIXED,
+                                                   1.0f,
+                                                   false, &result.positions);
 
             std::vector<float> axis_thetas = {
                 static_cast<float>(theta) * calc_ntk_factor(t_extrapolation_ratio, axes_dim[0]),
                 static_cast<float>(theta) * calc_ntk_factor(h_extrapolation_ratio, axes_dim[1]),
                 static_cast<float>(theta) * calc_ntk_factor(w_extrapolation_ratio, axes_dim[2]),
             };
-            return Rope::embed_nd(ids, bs, axis_thetas, axes_dim);
+            result.values = Rope::embed_nd(result.ids, bs, axis_thetas, axes_dim, result.layout, &result.frequencies);
+            return result;
         }
 
         ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
@@ -657,16 +660,16 @@ namespace Anima {
             int64_t h_pad = x->ne[1] + pad_h;
             int64_t w_pad = x->ne[0] + pad_w;
 
-            image_pe_vec          = gen_anima_image_pe_vec(1,
-                                                           static_cast<int>(h_pad),
-                                                           static_cast<int>(w_pad),
-                                                           static_cast<int>(config.patch_size),
-                                                           config.theta,
-                                                           config.axes_dim,
-                                                           4.0f,
-                                                           4.0f,
-                                                           1.0f,
-                                                           ref_latents);
+            image_pe_vec          = finish_rope_pe(gen_anima_image_pe_vec(1,
+                                                                          static_cast<int>(h_pad),
+                                                                          static_cast<int>(w_pad),
+                                                                          static_cast<int>(config.patch_size),
+                                                                          config.theta,
+                                                                          config.axes_dim,
+                                                                          4.0f,
+                                                                          4.0f,
+                                                                          1.0f,
+                                                                          ref_latents));
             int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2));
             auto image_pe         = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, image_pos_len);
             set_backend_tensor_data(image_pe, image_pe_vec.data());

src/model/diffusion/boogu.hpp

@@ -720,15 +720,18 @@ namespace Boogu {
         }
     }
 
-    __STATIC_INLINE__ std::vector<float> gen_boogu_pe(int h,
-                                                      int w,
-                                                      int patch_size,
-                                                      int bs,
-                                                      int context_len,
-                                                      const std::vector<ggml_tensor*>& ref_latents,
-                                                      int theta,
-                                                      const std::vector<int>& axes_dim) {
-        std::vector<std::vector<float>> ids;
+    __STATIC_INLINE__ Rope::Embedding gen_boogu_pe(int h,
+                                                   int w,
+                                                   int patch_size,
+                                                   int bs,
+                                                   int context_len,
+                                                   const std::vector<ggml_tensor*>& ref_latents,
+                                                   int theta,
+                                                   const std::vector<int>& axes_dim) {
+        Rope::Embedding result;
+        result.batch_size = bs;
+        result.positions.append_tokens(context_len);
+        auto& ids = result.ids;
         ids.reserve(static_cast<size_t>(bs) * context_len);
         for (int b = 0; b < bs; b++) {
             for (int i = 0; i < context_len; i++) {
@@ -741,15 +744,18 @@ namespace Boogu {
         for (ggml_tensor* ref : ref_latents) {
             int ref_h_tokens = patched_token_count(ref->ne[1], patch_size);
             int ref_w_tokens = patched_token_count(ref->ne[0], patch_size);
+            result.positions.append_image(ref_h_tokens, ref_w_tokens);
             append_spatial_ids(ids, bs, pe_shift, ref_h_tokens, ref_w_tokens);
             pe_shift += std::max(ref_h_tokens, ref_w_tokens);
         }
 
         int h_tokens = patched_token_count(h, patch_size);
         int w_tokens = patched_token_count(w, patch_size);
+        result.positions.append_image(h_tokens, w_tokens);
         append_spatial_ids(ids, bs, pe_shift, h_tokens, w_tokens);
 
-        return Rope::embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
+        result.values = Rope::embed_nd(ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
+        return result;
     }
 
     struct BooguImageRunner : public DiffusionModelRunner {
@@ -793,14 +799,14 @@ namespace Boogu {
                 ref_latents.push_back(make_input(ref_latent_tensor));
             }
 
-            pe_vec      = gen_boogu_pe(static_cast<int>(x->ne[1]),
-                                       static_cast<int>(x->ne[0]),
-                                       config.patch_size,
-                                       static_cast<int>(x->ne[3]),
-                                       static_cast<int>(context->ne[1]),
-                                       ref_latents,
-                                       config.theta,
-                                       config.axes_dim);
+            pe_vec      = finish_rope_pe(gen_boogu_pe(static_cast<int>(x->ne[1]),
+                                                      static_cast<int>(x->ne[0]),
+                                                      config.patch_size,
+                                                      static_cast<int>(x->ne[3]),
+                                                      static_cast<int>(context->ne[1]),
+                                                      ref_latents,
+                                                      config.theta,
+                                                      config.axes_dim));
             int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
             auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
             set_backend_tensor_data(pe, pe_vec.data());

src/model/diffusion/ernie_image.hpp

@@ -415,15 +415,13 @@ namespace ErnieImage {
             GGML_ASSERT(!context_tensor.empty());
             ggml_tensor* context = make_input(context_tensor);
 
-            pe_vec      = Rope::gen_ernie_image_pe(static_cast<int>(x->ne[1]),
-                                                   static_cast<int>(x->ne[0]),
-                                                   config.patch_size,
-                                                   static_cast<int>(x->ne[3]),
-                                                   static_cast<int>(context->ne[1]),
-                                                   config.theta,
-                                                   circular_y_enabled,
-                                                   circular_x_enabled,
-                                                   config.axes_dim);
+            pe_vec      = finish_rope_pe(Rope::gen_ernie_image_pe(static_cast<int>(x->ne[1]),
+                                                                  static_cast<int>(x->ne[0]),
+                                                                  config.patch_size,
+                                                                  static_cast<int>(x->ne[3]),
+                                                                  static_cast<int>(context->ne[1]),
+                                                                  config.theta,
+                                                                  config.axes_dim));
             int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
             auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, config.axes_dim_sum, 1, pos_len, 2);
             set_backend_tensor_data(pe, pe_vec.data());

src/model/diffusion/flux.hpp

@@ -1548,20 +1548,18 @@ namespace Flux {
             } else if (version == VERSION_OVIS_IMAGE) {
                 txt_arange_dims = {1, 2};
             }
-            pe_vec      = Rope::gen_flux_pe(static_cast<int>(x->ne[1]),
-                                            static_cast<int>(x->ne[0]),
-                                            config.patch_size,
-                                            static_cast<int>(x->ne[3]),
-                                            static_cast<int>(context->ne[1]),
-                                            txt_arange_dims,
-                                            ref_latents,
-                                            ref_index_mode,
-                                            config.ref_index_scale,
-                                            config.theta,
-                                            circular_y_enabled,
-                                            circular_x_enabled,
-                                            config.axes_dim,
-                                            sd_version_is_longcat(version));
+            pe_vec      = finish_rope_pe(Rope::gen_flux_pe(static_cast<int>(x->ne[1]),
+                                                           static_cast<int>(x->ne[0]),
+                                                           config.patch_size,
+                                                           static_cast<int>(x->ne[3]),
+                                                           static_cast<int>(context->ne[1]),
+                                                           txt_arange_dims,
+                                                           ref_latents,
+                                                           ref_index_mode,
+                                                           config.ref_index_scale,
+                                                           config.theta,
+                                                           config.axes_dim,
+                                                           sd_version_is_longcat(version)));
             int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
             // LOG_VERBOSE("pos_len %d", pos_len);
             auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);

src/model/diffusion/ideogram4.hpp

@@ -149,18 +149,21 @@ namespace Ideogram4 {
         return std::make_shared<Linear>(in_features, out_features, bias);
     }
 
-    __STATIC_INLINE__ std::vector<float> gen_ideogram4_pe(int grid_h,
-                                                          int grid_w,
-                                                          int bs,
-                                                          int context_len,
-                                                          int head_dim,
-                                                          int rope_theta,
-                                                          const std::vector<int>& mrope_section,
-                                                          bool circular_x = false,
-                                                          bool circular_y = false) {
+    __STATIC_INLINE__ Rope::Embedding gen_ideogram4_pe(int grid_h,
+                                                       int grid_w,
+                                                       int bs,
+                                                       int context_len,
+                                                       int head_dim,
+                                                       int rope_theta,
+                                                       const std::vector<int>& mrope_section) {
         GGML_ASSERT(bs == 1);
-        std::vector<std::vector<float>> ids(static_cast<size_t>(bs) * (context_len + grid_h * grid_w),
-                                            std::vector<float>(3, 0.f));
+        Rope::Embedding result;
+        result.batch_size = bs;
+        result.positions.append_tokens(context_len);
+        result.positions.append_image(grid_h, grid_w);
+        result.ids.assign(static_cast<size_t>(bs) * (context_len + grid_h * grid_w),
+                          std::vector<float>(3, 0.f));
+        auto& ids = result.ids;
 
         for (int i = 0; i < context_len; ++i) {
             ids[i] = {static_cast<float>(i), static_cast<float>(i), static_cast<float>(i)};
@@ -175,29 +178,13 @@ namespace Ideogram4 {
             }
         }
 
-        std::vector<std::vector<int>> axis_wrap_dims(3);
-        if (circular_y || circular_x) {
-            size_t total_len = static_cast<size_t>(bs) * (context_len + grid_h * grid_w);
-            axis_wrap_dims[1].assign(total_len, 0);
-            axis_wrap_dims[2].assign(total_len, 0);
-            if (circular_y) {
-                for (size_t idx = static_cast<size_t>(context_len); idx < total_len; ++idx) {
-                    axis_wrap_dims[1][idx] = grid_h;
-                }
-            }
-            if (circular_x) {
-                for (size_t idx = static_cast<size_t>(context_len); idx < total_len; ++idx) {
-                    axis_wrap_dims[2][idx] = grid_w;
-                }
-            }
-        }
-
-        return Rope::embed_interleaved_mrope(ids,
-                                             bs,
-                                             static_cast<float>(rope_theta),
-                                             head_dim,
-                                             mrope_section,
-                                             axis_wrap_dims);
+        result.values = Rope::embed_interleaved_mrope(ids,
+                                                      bs,
+                                                      static_cast<float>(rope_theta),
+                                                      head_dim,
+                                                      mrope_section,
+                                                      &result.frequencies);
+        return result;
     }
 
     class Ideogram4Attention : public GGMLBlock {
@@ -509,15 +496,13 @@ namespace Ideogram4 {
             int64_t head_dim = config.emb_dim / config.num_heads;
 
             auto runner_ctx = get_context();
-            pe_vec          = gen_ideogram4_pe(static_cast<int>(grid_h),
-                                               static_cast<int>(grid_w),
-                                               static_cast<int>(x->ne[3]),
-                                               static_cast<int>(context_len),
-                                               static_cast<int>(head_dim),
-                                               static_cast<int>(config.rope_theta),
-                                               config.mrope_section,
-                                               runner_ctx.circular_x_enabled,
-                                               runner_ctx.circular_y_enabled);
+            pe_vec          = finish_rope_pe(gen_ideogram4_pe(static_cast<int>(grid_h),
+                                                              static_cast<int>(grid_w),
+                                                              static_cast<int>(x->ne[3]),
+                                                              static_cast<int>(context_len),
+                                                              static_cast<int>(head_dim),
+                                                              static_cast<int>(config.rope_theta),
+                                                              config.mrope_section));
             auto pe         = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len);
             set_backend_tensor_data(pe, pe_vec.data());
 

src/model/diffusion/krea2.hpp

@@ -689,23 +689,28 @@ namespace Krea2 {
         }
     };
 
-    __STATIC_INLINE__ std::vector<float> gen_krea2_pe(int h,
-                                                      int w,
-                                                      int patch_size,
-                                                      int bs,
-                                                      int context_len,
-                                                      float theta,
-                                                      const std::vector<int>& axes_dim,
-                                                      const std::vector<ggml_tensor*>& ref_latents,
-                                                      Rope::RefIndexMode ref_index_mode) {
+    __STATIC_INLINE__ Rope::Embedding gen_krea2_pe(int h,
+                                                   int w,
+                                                   int patch_size,
+                                                   int bs,
+                                                   int context_len,
+                                                   float theta,
+                                                   const std::vector<int>& axes_dim,
+                                                   const std::vector<ggml_tensor*>& ref_latents,
+                                                   Rope::RefIndexMode ref_index_mode) {
+        Rope::Embedding result;
+        result.batch_size = bs;
+        result.positions.append_tokens(context_len);
         auto txt_ids = Rope::gen_flux_txt_ids(bs, context_len, 3, {});
-        auto img_ids = Rope::gen_flux_img_ids(h, w, patch_size, bs, 3, 0, 0, 0, false);
+        auto img_ids = Rope::gen_flux_img_ids(h, w, patch_size, bs, 3, 0, 0, 0, false, &result.positions);
         auto ids     = Rope::concat_ids(txt_ids, img_ids, bs);
         if (ref_latents.size() > 0) {
-            auto refs_ids = Rope::gen_refs_ids(patch_size, bs, 3, 1, ref_latents, ref_index_mode, 1.0f, false, 0);
+            auto refs_ids = Rope::gen_refs_ids(patch_size, bs, 3, 1, ref_latents, ref_index_mode, 1.0f, false, 0, &result.positions);
             ids           = Rope::concat_ids(ids, refs_ids, bs);
         }
-        return Rope::embed_nd(ids, bs, theta, axes_dim);
+        result.ids    = std::move(ids);
+        result.values = Rope::embed_nd(result.ids, bs, theta, axes_dim, result.layout, &result.frequencies);
+        return result;
     }
 
     struct Krea2Runner : public DiffusionModelRunner {
@@ -749,15 +754,15 @@ namespace Krea2 {
                 ref_latents.push_back(make_input(ref_latent_tensor));
             }
 
-            pe_vec      = gen_krea2_pe(static_cast<int>(x->ne[1]),
-                                       static_cast<int>(x->ne[0]),
-                                       config.patch_size,
-                                       static_cast<int>(x->ne[3]),
-                                       static_cast<int>(context->ne[1]),
-                                       config.theta,
-                                       config.axes_dim,
-                                       ref_latents,
-                                       ref_image_params.ref_index_mode);
+            pe_vec      = finish_rope_pe(gen_krea2_pe(static_cast<int>(x->ne[1]),
+                                                      static_cast<int>(x->ne[0]),
+                                                      config.patch_size,
+                                                      static_cast<int>(x->ne[3]),
+                                                      static_cast<int>(context->ne[1]),
+                                                      config.theta,
+                                                      config.axes_dim,
+                                                      ref_latents,
+                                                      ref_image_params.ref_index_mode));
             int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
             auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
             set_backend_tensor_data(pe, pe_vec.data());

src/model/diffusion/lens.hpp

@@ -384,14 +384,12 @@ namespace Lens {
             GGML_ASSERT(!context_tensor.empty());
             ggml_tensor* context = make_input(context_tensor);
 
-            pe_vec      = Rope::gen_lens_pe(static_cast<int>(x->ne[1]),
-                                            static_cast<int>(x->ne[0]),
-                                            static_cast<int>(x->ne[3]),
-                                            static_cast<int>(context->ne[1]),
-                                            config.theta,
-                                            circular_y_enabled,
-                                            circular_x_enabled,
-                                            config.axes_dim);
+            pe_vec      = finish_rope_pe(Rope::gen_lens_pe(static_cast<int>(x->ne[1]),
+                                                           static_cast<int>(x->ne[0]),
+                                                           static_cast<int>(x->ne[3]),
+                                                           static_cast<int>(context->ne[1]),
+                                                           config.theta,
+                                                           config.axes_dim));
             int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
             auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
             set_backend_tensor_data(pe, pe_vec.data());

src/model/diffusion/llada_image.hpp

@@ -412,16 +412,14 @@ namespace LLaDAImage {
             GGML_ASSERT(!context_tensor.empty());
             ggml_tensor* context = make_input(context_tensor);
 
-            pe_vec      = Rope::gen_llada_image_pe(static_cast<int>(x->ne[1]),
-                                                   static_cast<int>(x->ne[0]),
-                                                   config.patch_size,
-                                                   static_cast<int>(x->ne[3]),
-                                                   static_cast<int>(context->ne[1]),
-                                                   ZImage::SEQ_MULTI_OF,
-                                                   config.theta,
-                                                   circular_y_enabled,
-                                                   circular_x_enabled,
-                                                   config.axes_dim);
+            pe_vec      = finish_rope_pe(Rope::gen_llada_image_pe(static_cast<int>(x->ne[1]),
+                                                                  static_cast<int>(x->ne[0]),
+                                                                  config.patch_size,
+                                                                  static_cast<int>(x->ne[3]),
+                                                                  static_cast<int>(context->ne[1]),
+                                                                  ZImage::SEQ_MULTI_OF,
+                                                                  config.theta,
+                                                                  config.axes_dim));
             int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
             auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
             set_backend_tensor_data(pe, pe_vec.data());
@@ -461,14 +459,14 @@ namespace LLaDAImage {
             ggml_tensor* source    = make_input(source_tensor);
             GGML_ASSERT(x->ne[3] == 1);
 
-            pe_vec      = Rope::gen_llada_image_edit_pe(static_cast<int>(x->ne[1]),
-                                                        static_cast<int>(x->ne[0]),
-                                                        config.patch_size,
-                                                        static_cast<int>(context->ne[1]),
-                                                   semantic != nullptr ? static_cast<int>(semantic->ne[1]) : 0,
-                                                        ZImage::SEQ_MULTI_OF,
-                                                        config.theta,
-                                                        config.axes_dim);
+            pe_vec      = finish_rope_pe(Rope::gen_llada_image_edit_pe(static_cast<int>(x->ne[1]),
+                                                                       static_cast<int>(x->ne[0]),
+                                                                       config.patch_size,
+                                                                       static_cast<int>(context->ne[1]),
+                                                                  semantic != nullptr ? static_cast<int>(semantic->ne[1]) : 0,
+                                                                       ZImage::SEQ_MULTI_OF,
+                                                                       config.theta,
+                                                                       config.axes_dim));
             int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
             auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
             set_backend_tensor_data(pe, pe_vec.data());

src/model/diffusion/mage_flow.hpp

@@ -110,13 +110,13 @@ namespace MageFlow {
             }
 
             int batch_size = static_cast<int>(x->ne[3]);
-            pe_vec         = Rope::gen_mage_flow_pe(static_cast<int>(x->ne[1]),
-                                                    static_cast<int>(x->ne[0]),
-                                                    batch_size,
-                                                    static_cast<int>(context->ne[1]),
-                                                    ref_latents,
-                                                    config.theta,
-                                                    config.axes_dim);
+            pe_vec         = finish_rope_pe(Rope::gen_mage_flow_pe(static_cast<int>(x->ne[1]),
+                                                                   static_cast<int>(x->ne[0]),
+                                                                   batch_size,
+                                                                   static_cast<int>(context->ne[1]),
+                                                                   ref_latents,
+                                                                   config.theta,
+                                                                   config.axes_dim));
             int pos_len    = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
             auto pe        = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
             set_backend_tensor_data(pe, pe_vec.data());

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/minit2i.hpp

@@ -154,18 +154,26 @@ namespace MiniT2I {
         return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), head_dim, 10000.f));
     }
 
-    inline std::vector<float> make_vision_rope(int side, int head_dim) {
+    inline Rope::Embedding make_vision_rope(int side, int head_dim) {
         GGML_ASSERT(head_dim % 4 == 0);
         int dim     = head_dim / 2;
         int quarter = dim / 2;
         int length  = side * side;
+        Rope::Embedding result;
+        result.positions.append_image(side, side);
         std::vector<float> out(static_cast<size_t>(length) * (head_dim / 2) * 4);
         std::vector<float> freqs(quarter);
         for (int i = 0; i < quarter; ++i) {
             freqs[i] = 1.0f / std::pow(10000.0f, static_cast<float>(2 * i) / static_cast<float>(dim));
         }
+        for (int axis : {1, 2}) {
+            for (float frequency : freqs) {
+                result.frequencies.push_back({static_cast<size_t>(axis), frequency});
+            }
+        }
         for (int y = 0; y < side; ++y) {
             for (int x = 0; x < side; ++x) {
+                result.ids.push_back({0.f, static_cast<float>(y), static_cast<float>(x)});
                 int pos     = y * side + x;
                 size_t base = static_cast<size_t>(pos) * (head_dim / 2) * 4;
                 for (int i = 0; i < quarter; ++i) {
@@ -182,7 +190,8 @@ namespace MiniT2I {
                 }
             }
         }
-        return out;
+        result.values = std::move(out);
+        return result;
     }
 
     struct SwiGLUMlp : public GGMLBlock {
@@ -475,6 +484,8 @@ namespace MiniT2I {
         int64_t cached_txt_len                      = -1;
         int64_t cached_hidden_size                  = -1;
         int64_t cached_head_dim                     = -1;
+        bool cached_circular_x                      = false;
+        bool cached_circular_y                      = false;
 
         MiniT2IRunner(ggml_backend_t backend,
                       const String2TensorStorage& tensor_storage_map      = {},
@@ -521,6 +532,8 @@ namespace MiniT2I {
                 cached_txt_len == txt_len &&
                 cached_hidden_size == config.hidden_size &&
                 cached_head_dim == config.head_dim &&
+                cached_circular_x == circular_x_enabled &&
+                cached_circular_y == circular_y_enabled &&
                 cached_pos_embed != nullptr &&
                 cached_txt_pe != nullptr &&
                 cached_joint_pe != nullptr) {
@@ -531,7 +544,7 @@ namespace MiniT2I {
 
             auto pos_embed_vec = make_2d_sincos_pos_embed(static_cast<int>(img_side), static_cast<int>(config.hidden_size));
             auto txt_pe_vec    = make_text_rope(static_cast<int>(txt_len), static_cast<int>(config.head_dim));
-            auto img_pe_vec    = make_vision_rope(static_cast<int>(img_side), static_cast<int>(config.head_dim));
+            auto img_pe_vec    = finish_rope_pe(make_vision_rope(static_cast<int>(img_side), static_cast<int>(config.head_dim)));
             auto joint_pe_vec  = txt_pe_vec;
             joint_pe_vec.insert(joint_pe_vec.end(), img_pe_vec.begin(), img_pe_vec.end());
 
@@ -561,6 +574,8 @@ namespace MiniT2I {
             cached_txt_len     = txt_len;
             cached_hidden_size = config.hidden_size;
             cached_head_dim    = config.head_dim;
+            cached_circular_x  = circular_x_enabled;
+            cached_circular_y  = circular_y_enabled;
         }
 
         ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,

src/model/diffusion/model.hpp

@@ -7,7 +7,7 @@
 
 #include "core/ggml_runner.h"
 #include "core/tensor_ggml.hpp"
-#include "model/common/rope.hpp"
+#include "model/common/rope_circular.hpp"
 #include "model_manager.h"
 
 enum class RefImageResizeMode {
@@ -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 {
@@ -182,6 +184,11 @@ struct DiffusionModelRunner : public GGMLRunner {
 protected:
     std::string prefix;
 
+    std::vector<float> finish_rope_pe(Rope::Embedding embedding) {
+        Rope::apply_circular(embedding, circular_x_enabled, circular_y_enabled);
+        return std::move(embedding.values);
+    }
+
 public:
     DiffusionModelRunner(ggml_backend_t backend,
                          const std::string& prefix,

src/model/diffusion/pid.hpp

@@ -135,13 +135,13 @@ namespace Pid {
         return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), dim, theta));
     }
 
-    inline std::vector<float> make_rope_2d(int height,
-                                           int width,
-                                           int dim,
-                                           float theta    = 10000.f,
-                                           float scale    = 16.f,
-                                           int ref_grid_h = 0,
-                                           int ref_grid_w = 0) {
+    inline Rope::Embedding make_rope_2d(int height,
+                                        int width,
+                                        int dim,
+                                        float theta    = 10000.f,
+                                        float scale    = 16.f,
+                                        int ref_grid_h = 0,
+                                        int ref_grid_w = 0) {
         GGML_ASSERT(dim % 4 == 0);
         return Rope::embed_2d_interleaved(height, width, dim, theta, scale, ref_grid_h, ref_grid_w);
     }
@@ -867,13 +867,13 @@ namespace Pid {
             int64_t Hs = Hp / config.patch_size;
             int64_t Ws = Wp / config.patch_size;
 
-            pos_img_vec  = make_rope_2d(static_cast<int>(Hs),
-                                        static_cast<int>(Ws),
-                                        static_cast<int>(config.hidden_size / config.num_groups),
-                                        10000.f,
-                                        16.f,
-                                        static_cast<int>(config.rope_ref_grid_h),
-                                        static_cast<int>(config.rope_ref_grid_w));
+            pos_img_vec  = finish_rope_pe(make_rope_2d(static_cast<int>(Hs),
+                                                       static_cast<int>(Ws),
+                                                       static_cast<int>(config.hidden_size / config.num_groups),
+                                                       10000.f,
+                                                       16.f,
+                                                       static_cast<int>(config.rope_ref_grid_h),
+                                                       static_cast<int>(config.rope_ref_grid_w)));
             auto pos_img = ggml_new_tensor_4d(compute_ctx,
                                               GGML_TYPE_F32,
                                               2,
@@ -904,13 +904,13 @@ namespace Pid {
                                                 1);
             set_backend_tensor_data(pixel_pos, pixel_pos_vec.data());
 
-            pixel_pos_comp_vec  = make_rope_2d(static_cast<int>(Hs),
-                                               static_cast<int>(Ws),
-                                               static_cast<int>(config.pixel_attn_hidden_size / config.pixel_num_groups),
-                                               10000.f,
-                                               16.f,
-                                               static_cast<int>(config.rope_ref_grid_h),
-                                               static_cast<int>(config.rope_ref_grid_w));
+            pixel_pos_comp_vec  = finish_rope_pe(make_rope_2d(static_cast<int>(Hs),
+                                                              static_cast<int>(Ws),
+                                                              static_cast<int>(config.pixel_attn_hidden_size / config.pixel_num_groups),
+                                                              10000.f,
+                                                              16.f,
+                                                              static_cast<int>(config.rope_ref_grid_h),
+                                                              static_cast<int>(config.rope_ref_grid_w)));
             auto pixel_pos_comp = ggml_new_tensor_4d(compute_ctx,
                                                      GGML_TYPE_F32,
                                                      2,

src/model/diffusion/qwen_image.hpp

@@ -635,18 +635,16 @@ namespace Qwen {
                 ref_index_mode = Rope::RefIndexMode::DECREASE;
             }
 
-            pe_vec      = Rope::gen_qwen_image_pe(time_len,
-                                                  static_cast<int>(x->ne[1]),
-                                                  static_cast<int>(x->ne[0]),
-                                                  config.patch_size,
-                                                  batch_size,
-                                                  static_cast<int>(context->ne[1]),
-                                                  ref_latents,
-                                                  ref_index_mode,
-                                                  config.theta,
-                                                  circular_y_enabled,
-                                                  circular_x_enabled,
-                                                  config.axes_dim);
+            pe_vec      = finish_rope_pe(Rope::gen_qwen_image_pe(time_len,
+                                                                 static_cast<int>(x->ne[1]),
+                                                                 static_cast<int>(x->ne[0]),
+                                                                 config.patch_size,
+                                                                 batch_size,
+                                                                 static_cast<int>(context->ne[1]),
+                                                                 ref_latents,
+                                                                 ref_index_mode,
+                                                                 config.theta,
+                                                                 config.axes_dim));
             int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
             // LOG_VERBOSE("pos_len %d", pos_len);
             auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);

src/model/diffusion/qwen_image_2_1.hpp

@@ -68,6 +68,7 @@ namespace Qwen {
         std::vector<QwenImage21Segment> segments;
         std::vector<std::vector<float>> positions;
         int64_t prefix_length = 0;
+        Rope::PositionLayout rope_layout;
 
         static QwenImage21Layout build(int64_t text_length,
                                        const sd::Tensor<int32_t>& image_slots,
@@ -82,6 +83,7 @@ namespace Qwen {
                 auto [height, width] = image_shapes[index];
                 int64_t start        = static_cast<int64_t>(layout.positions.size());
                 layout.segments.push_back({start, start + height * width, context_start, index});
+                layout.rope_layout.append_image(static_cast<int>(height), static_cast<int>(width));
                 for (int64_t h = 0; h < height; ++h) {
                     for (int64_t w = 0; w < width; ++w) {
                         layout.positions.push_back({static_cast<float>(position),
@@ -106,6 +108,7 @@ namespace Qwen {
                 } else {
                     int64_t start = static_cast<int64_t>(layout.positions.size());
                     layout.segments.push_back({start, start + i - begin, begin, -1});
+                    layout.rope_layout.append_tokens(i - begin);
                     for (int64_t j = begin; j < i; ++j, ++position) {
                         float p = static_cast<float>(position);
                         layout.positions.push_back({p, p, p});
@@ -121,6 +124,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 +175,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 +256,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 +277,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 +299,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 +307,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 +351,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 +372,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 +416,87 @@ 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) +
+                             ".circular." + std::to_string(circular_x_enabled) + std::to_string(circular_y_enabled);
+                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);
+                Rope::Embedding embedding;
+                embedding.ids.assign(first_position, layout.positions.end());
+                const size_t offset             = cached ? static_cast<size_t>(layout.prefix_length) : 0;
+                embedding.positions.token_count = embedding.ids.size();
+                for (auto region : layout.rope_layout.images) {
+                    if (region.begin >= offset) {
+                        region.begin -= offset;
+                        embedding.positions.images.push_back(region);
                     }
                 }
-                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));
+                embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, config.axes_dim, embedding.layout, &embedding.frequencies);
+                pe_data          = finish_rope_pe(std::move(embedding));
+                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));
+                    }
                 }
-                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/diffusion/z_image.hpp

@@ -642,18 +642,16 @@ namespace ZImage {
                 ref_latents.push_back(make_input(ref_latent_tensor));
             }
 
-            pe_vec      = Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
-                                               static_cast<int>(x->ne[0]),
-                                               config.patch_size,
-                                               static_cast<int>(x->ne[3]),
-                                               static_cast<int>(context->ne[1]),
-                                               SEQ_MULTI_OF,
-                                               ref_latents,
-                                               ref_index_mode,
-                                               config.theta,
-                                               circular_y_enabled,
-                                               circular_x_enabled,
-                                               config.axes_dim);
+            pe_vec      = finish_rope_pe(Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
+                                                              static_cast<int>(x->ne[0]),
+                                                              config.patch_size,
+                                                              static_cast<int>(x->ne[3]),
+                                                              static_cast<int>(context->ne[1]),
+                                                              SEQ_MULTI_OF,
+                                                              ref_latents,
+                                                              ref_index_mode,
+                                                              config.theta,
+                                                              config.axes_dim));
             int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
             // LOG_VERBOSE("pos_len %d", pos_len);
             auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);

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/model_loader.cpp

@@ -874,7 +874,8 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
 
 std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
                                                        std::set<std::string> ignore_tensors,
-                                                       bool writable_mmap) {
+                                                       bool writable_mmap,
+                                                       ggml_backend_dev_t device) {
     std::set<std::string> names;
     for (const auto& entry : tensors) {
         names.insert(entry.first);
@@ -896,6 +897,39 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
         if (!fdata.mmbuffer)
             continue;
 
+        // Wrapped on first use: a device buffer makes the whole file resident on that device.
+        std::shared_ptr<struct ggml_backend_buffer> file_buffer = device == nullptr ? fdata.mmbuffer : nullptr;
+        bool file_unmappable                                    = false;
+
+        auto buffer_for_file = [&]() -> ggml_backend_buffer_t {
+            if (file_buffer || file_unmappable) {
+                return file_buffer.get();
+            }
+            auto cached = fdata.device_mmbuffers.find(device);
+            if (cached != fdata.device_mmbuffers.end()) {
+                file_buffer = cached->second;
+                return file_buffer.get();
+            }
+            size_t max_tensor_size = 0;
+            for (const auto& ts : fdata.tensors) {
+                max_tensor_size = std::max(max_tensor_size, static_cast<size_t>(ts.nbytes()));
+            }
+            ggml_backend_buffer_t buf = sd_backend_dev_buffer_from_host_ptr(device,
+                                                                            fdata.mmapped->writable_data(),
+                                                                            fdata.mmapped->size(),
+                                                                            max_tensor_size);
+            if (buf == nullptr) {
+                LOG_WARN("mmap: %s cannot map '%s', loading it instead",
+                         ggml_backend_dev_name(device), fdata.path.c_str());
+                file_unmappable = true;
+                return nullptr;
+            }
+            LOG_INFO("mmap: mapped '%s' for %s", fdata.path.c_str(), ggml_backend_dev_name(device));
+            file_buffer                    = std::shared_ptr<struct ggml_backend_buffer>(buf, ggml_backend_buffer_free);
+            fdata.device_mmbuffers[device] = file_buffer;
+            return file_buffer.get();
+        };
+
         const std::vector<TensorStorage>& file_tensors = fdata.tensors;
 
         size_t file_mapped_bytes   = 0;
@@ -944,10 +978,13 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
                 continue;
             }
 
-            ggml_backend_buffer_t buf_mmap = fdata.mmbuffer.get();
-            uint8_t* mmap_data             = static_cast<uint8_t*>(ggml_backend_buffer_get_base(buf_mmap));
-            dst_tensor->buffer             = buf_mmap;
-            dst_tensor->data               = mmap_data + tensor_offset;
+            ggml_backend_buffer_t buf_mmap = buffer_for_file();
+            if (buf_mmap == nullptr) {
+                break;
+            }
+            uint8_t* mmap_data = static_cast<uint8_t*>(ggml_backend_buffer_get_base(buf_mmap));
+            dst_tensor->buffer = buf_mmap;
+            dst_tensor->data   = mmap_data + tensor_offset;
 
             file_mapped_bytes += tensor_size;
             file_mapped_tensors++;
@@ -956,7 +993,7 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
         if (file_mapped_bytes > 0) {
             mapped_tensors += file_mapped_tensors;
             mapped_bytes += file_mapped_bytes;
-            result.push_back({fdata.mmapped, fdata.mmbuffer});
+            result.push_back({fdata.mmapped, file_buffer});
         }
     }
 
@@ -972,6 +1009,16 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
     return result;
 }
 
+std::vector<ggml_backend_buffer_t> ModelLoader::get_device_mmap_buffers() const {
+    std::vector<ggml_backend_buffer_t> buffers;
+    for (const auto& fdata : file_data) {
+        for (const auto& entry : fdata.device_mmbuffers) {
+            buffers.push_back(entry.second.get());
+        }
+    }
+    return buffers;
+}
+
 bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                                bool enable_mmap,
                                const std::set<std::string>* target_tensor_names,
@@ -1115,6 +1162,11 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                     if (dst_tensor->buffer != nullptr && dst_tensor->buffer == fdata.mmbuffer.get()) {
                         continue;
                     }
+                    if (dst_tensor->buffer != nullptr &&
+                        std::any_of(fdata.device_mmbuffers.begin(), fdata.device_mmbuffers.end(),
+                                    [&](const auto& entry) { return entry.second.get() == dst_tensor->buffer; })) {
+                        continue;
+                    }
 
                     size_t nbytes_to_read = tensor_storage.nbytes_to_read();
 

src/model_loader.h

@@ -20,6 +20,8 @@ struct ModelFileData {
     std::vector<TensorStorage> tensors;
     std::shared_ptr<MmapWrapper> mmapped;
     std::shared_ptr<struct ggml_backend_buffer> mmbuffer;
+    // mmapped wrapped by devices that can use host memory in place (buffer_from_host_ptr)
+    std::map<ggml_backend_dev_t, std::shared_ptr<struct ggml_backend_buffer>> device_mmbuffers;
     bool is_zip;
 };
 
@@ -120,7 +122,9 @@ public:
     void process_model_files(bool enable_mmap = false, bool writable_mmap = true);
     std::vector<MmapTensorStore> mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
                                               std::set<std::string> ignore_tensors = {},
-                                              bool writable                        = true);
+                                              bool writable                        = true,
+                                              ggml_backend_dev_t device            = nullptr);
+    std::vector<ggml_backend_buffer_t> get_device_mmap_buffers() const;
     bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                       bool use_mmap                                    = false,
                       const std::set<std::string>* target_tensor_names = nullptr,

src/model_manager.cpp

@@ -780,38 +780,52 @@ bool ModelManager::validate_tensor(const TensorState& state) const {
 
 bool ModelManager::mmap_params(const std::vector<TensorState*>& states,
                                std::vector<ParamsStorageBlock*>& created_storage_blocks) {
-    std::map<std::string, ggml_tensor*> mmap_candidates;
-    std::map<std::string, TensorState*> mmap_states;
+    // A GPU that computes on mmapped params in place cannot address a CPU buffer, and nothing
+    // stages them for it, so they are mapped through a buffer of that GPU's device.
+    struct MmapGroup {
+        std::map<std::string, ggml_tensor*> candidates;
+        std::map<std::string, TensorState*> states;
+    };
+    std::map<ggml_backend_dev_t, MmapGroup> groups;
     for (TensorState* state : states) {
         if (state == nullptr || !can_mmap_storage(*state) || state->tensor == nullptr ||
             state->tensor->data != nullptr || state->tensor->view_src != nullptr) {
             continue;
         }
-        mmap_candidates[state->name] = state->tensor;
-        mmap_states[state->name]     = state;
-    }
-    if (mmap_candidates.empty()) {
-        return true;
-    }
-
-    auto mmap_store = model_loader_.mmap_tensors(mmap_candidates, {}, writable_mmap_);
-    if (mmap_store.empty()) {
-        return true;
-    }
-
-    auto block                = std::make_unique<ParamsStorageBlock>();
-    block->mmap_tensor_stores = std::move(mmap_store);
-    ParamsStorageBlock* raw   = block.get();
-    for (const auto& pair : mmap_states) {
-        TensorState* state = pair.second;
-        if (state != nullptr && state->tensor != nullptr && state->tensor->data != nullptr) {
-            block->states.push_back(state);
+        ggml_backend_dev_t device = nullptr;
+        if (!sd_backend_is_cpu(state->compute_backend) && !sd_backend_is_cpu(state->params_backend)) {
+            device = ggml_backend_get_device(state->compute_backend);
         }
+        MmapGroup& group              = groups[device];
+        group.candidates[state->name] = state->tensor;
+        group.states[state->name]     = state;
     }
 
-    if (!block->states.empty()) {
-        params_storage_blocks_.push_back(std::move(block));
-        created_storage_blocks.push_back(raw);
+    for (auto& [device, group] : groups) {
+        // Device buffers wrap read-only mappings only; params that LoRAs are merged into in place
+        // are loaded instead.
+        if (device != nullptr && writable_mmap_) {
+            continue;
+        }
+        auto mmap_store = model_loader_.mmap_tensors(group.candidates, {}, writable_mmap_, device);
+        if (mmap_store.empty()) {
+            continue;
+        }
+
+        auto block                = std::make_unique<ParamsStorageBlock>();
+        block->mmap_tensor_stores = std::move(mmap_store);
+        ParamsStorageBlock* raw   = block.get();
+        for (const auto& pair : group.states) {
+            TensorState* state = pair.second;
+            if (state != nullptr && state->tensor != nullptr && state->tensor->data != nullptr) {
+                block->states.push_back(state);
+            }
+        }
+
+        if (!block->states.empty()) {
+            params_storage_blocks_.push_back(std::move(block));
+            created_storage_blocks.push_back(raw);
+        }
     }
     return true;
 }
@@ -1353,15 +1367,16 @@ size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backe
     }
 
     size_t total_size = 0;
-    auto add_buffer   = [&](ggml_backend_buffer_t buffer) {
-        if (buffer == nullptr || ggml_backend_buffer_is_host(buffer)) {
+    std::unordered_set<ggml_backend_buffer_t> seen;
+    auto add_buffer = [&](ggml_backend_buffer_t buffer) {
+        if (buffer == nullptr || ggml_backend_buffer_is_host(buffer) || !seen.insert(buffer).second) {
             return;
         }
         ggml_backend_buffer_type_t buffer_type = ggml_backend_buffer_get_type(buffer);
         auto split_devices                     = split_buffer_devices_.find(buffer_type);
         const bool on_device                   = split_devices == split_buffer_devices_.end()
-                                                       ? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device
-                                                       : std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) {
+                                                     ? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device
+                                                     : std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) {
                                          return ggml_backend_get_device(entry.first) == compute_device;
                                      });
         if (!on_device) {
@@ -1371,9 +1386,16 @@ size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backe
         total_size               = buffer_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + buffer_size;
     };
 
+    // The loader may retain device mappings after their parameter blocks are released.
+    for (ggml_backend_buffer_t buffer : model_loader_.get_device_mmap_buffers()) {
+        add_buffer(buffer);
+    }
     for (const auto& block : params_storage_blocks_) {
         if (block != nullptr) {
             add_buffer(block->buffer);
+            for (const auto& store : block->mmap_tensor_stores) {
+                add_buffer(store.mmbuffer.get());
+            }
         }
     }
     for (const auto& block : compute_staging_blocks_) {

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,