perf: cache CPU position embeddings across DiT runners (#2211)

docs/pe_cache.md

@@ -0,0 +1,60 @@
+# Position embedding cache
+
+`sd::PECache` stores CPU-side `sd::Tensor<float>` data. Each model runner owns
+its cache and calls the existing `make_input()` on every graph build. Cache
+hits skip PE generation; allocation and upload remain ordinary graph-input
+operations. `GGMLRunner` has no PE-specific code or resource hooks.
+
+The cache keeps the most recent key and tensor for each named PE input. A
+matching key returns the existing tensor by reference. Changing its parameters
+replaces that input's cached result. Different named inputs have independent
+storage. Alternating two layouts under the same name recomputes their PE;
+there is no LRU, device-buffer cache, memory budget, or graph lifecycle callback.
+CPU data lives with the model runner, including across `runner_end()` calls,
+and is released when its named entry is replaced or the cache is destroyed.
+
+## Adding a model
+
+Include `core/pe_cache.h` and declare a member:
+
+```cpp
+sd::PECache pe_cache_;
+```
+
+Construct the key before generating PE and place the generation code in the
+lazy callback:
+
+```cpp
+sd::PECacheKey key("model.image_pe");
+key.add(height).add(width).add(head_dim).add(theta)
+    .add(circular_x_enabled).add(circular_y_enabled);
+auto pe = make_input(pe_cache_.get_or_create(key, [&]() {
+    auto values = generate_position_embedding(/* existing arguments */);
+    return sd::Tensor<float>({2, 2, head_dim / 2, token_count}, std::move(values));
+}));
+```
+
+Keys must cover every variable affecting PE values or shape: geometry, text
+length, frequency settings, reference-image ordering and offsets, circular
+settings, and video/audio positions or frame rates as applicable.
+`add_shape()` records all four ggml dimensions, including trailing singleton
+dimensions. `add_shapes()` records an ordered list of reference shapes.
+`add(sd::Tensor<T>)` records both shape and values for explicit position inputs.
+Keys compare their complete encoded fields rather than a data pointer or hash.
+
+Use distinct names for PE inputs that coexist in a graph, such as Anima's
+adapter query and key PE. `make_input()` defers reading the CPU data until graph
+execution: do not replace a named entry while an unfinished graph references
+it. Adding or updating other named entries does not invalidate its storage.
+Return tensor shapes with the contiguous dimension first, and preserve each
+model's existing F32 values, layout, and final transformations.
+
+## Current integrations
+
+The cache is used by MiniT2I, the Flux runner and its variants, Flux3, Qwen Image,
+Qwen Image 2.1, Z-Image and L2P, Anima, Boogu, ERNIE Image, Fibo2, Ideogram4,
+Iris, Krea2, Lens, LLaDA Image, MAGE Flow, Ming Image, MOSS Image, PiD, PixArt,
+Wan, Hunyuan Video, LingBot Video, and LTX audio/video.
+
+Position embeddings stored as model weights and graph-side RoPE operations
+remain on their existing paths, including MiniMax H3 and SenseNova U1.

src/core/pe_cache.cpp

@@ -0,0 +1,30 @@
+#include "core/pe_cache.h"
+
+namespace sd {
+    PECacheKey& PECacheKey::add_shape(const ggml_tensor* tensor) {
+        add(tensor != nullptr);
+        if (tensor != nullptr) {
+            for (int i = 0; i < GGML_MAX_DIMS; ++i) {
+                add(tensor->ne[i]);
+            }
+        }
+        return *this;
+    }
+
+    PECacheKey& PECacheKey::add_shapes(const std::vector<ggml_tensor*>& tensors) {
+        add(tensors.size());
+        for (const auto* tensor : tensors) {
+            add_shape(tensor);
+        }
+        return *this;
+    }
+
+    const Tensor<float>& PECache::get_or_create(const PECacheKey& key, const std::function<Tensor<float>()>& build) {
+        auto entry = entries_.find(key.name());
+        if (entry != entries_.end() && entry->second.key == key) {
+            return entry->second.value;
+        }
+        auto value = build();
+        return entries_.insert_or_assign(key.name(), Entry{key, std::move(value)}).first->second.value;
+    }
+}

src/core/pe_cache.h

@@ -0,0 +1,82 @@
+#ifndef __SD_CORE_PE_CACHE_H__
+#define __SD_CORE_PE_CACHE_H__
+
+#include <functional>
+#include <map>
+#include <set>
+#include <type_traits>
+
+#include "core/tensor.hpp"
+#include "ggml.h"
+
+namespace sd {
+    class PECacheKey {
+        std::string name_;
+        std::string fields_;
+
+    public:
+        explicit PECacheKey(std::string name)
+            : name_(std::move(name)) {}
+        const std::string& name() const { return name_; }
+
+        template <typename T, std::enable_if_t<std::is_arithmetic_v<T> || std::is_enum_v<T>, int> = 0>
+        PECacheKey& add(T value) {
+            if constexpr (std::is_floating_point_v<T>) {
+                const double field = static_cast<double>(value);
+                fields_.push_back('f');
+                fields_.append(reinterpret_cast<const char*>(&field), sizeof(field));
+            } else {
+                const uint64_t field = static_cast<uint64_t>(value);
+                fields_.push_back('i');
+                fields_.append(reinterpret_cast<const char*>(&field), sizeof(field));
+            }
+            return *this;
+        }
+
+        template <typename T>
+        PECacheKey& add(const std::vector<T>& values) {
+            fields_.push_back('v');
+            add(values.size());
+            for (const auto& value : values) {
+                add(value);
+            }
+            return *this;
+        }
+
+        template <typename T>
+        PECacheKey& add(const std::set<T>& values) {
+            return add(std::vector<T>(values.begin(), values.end()));
+        }
+
+        template <typename T>
+        PECacheKey& add(const Tensor<T>& values) {
+            add(values.shape());
+            for (int64_t i = 0; i < values.numel(); ++i) {
+                add(values.data()[i]);
+            }
+            return *this;
+        }
+
+        PECacheKey& add_shape(const ggml_tensor* tensor);
+        PECacheKey& add_shapes(const std::vector<ggml_tensor*>& tensors);
+
+        bool operator==(const PECacheKey& other) const {
+            return name_ == other.name_ && fields_ == other.fields_;
+        }
+    };
+
+    class PECache {
+        struct Entry {
+            PECacheKey key;
+            Tensor<float> value;
+        };
+
+        std::map<std::string, Entry> entries_;
+
+    public:
+        const Tensor<float>& get_or_create(const PECacheKey& key, const std::function<Tensor<float>()>& build);
+    };
+
+}
+
+#endif  // __SD_CORE_PE_CACHE_H__

src/model/diffusion/anima.hpp

@@ -7,6 +7,7 @@
 #include <memory>
 #include <utility>
 #include <vector>
+#include "core/pe_cache.h"
 
 #include "model/common/block.hpp"
 #include "model/common/rope.hpp"
@@ -562,10 +563,9 @@ namespace Anima {
     };
 
     struct AnimaRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
+
     public:
-        std::vector<float> image_pe_vec;
-        std::vector<float> adapter_q_pe_vec;
-        std::vector<float> adapter_k_pe_vec;
         AnimaConfig config;
         AnimaNet net;
 
@@ -660,7 +660,11 @@ namespace Anima {
             int64_t h_pad = x->ne[1] + pad_h;
             int64_t w_pad = x->ne[0] + pad_w;
 
-            image_pe_vec          = finish_rope_pe(gen_anima_image_pe_vec(1,
+            sd::PECacheKey image_pe_key("anima.image_pe");
+            image_pe_key.add(h_pad).add(w_pad).add(config.patch_size).add(config.theta).add(config.axes_dim).add(config.head_dim);
+            image_pe_key.add_shapes(ref_latents).add(circular_x_enabled).add(circular_y_enabled);
+            auto image_pe = make_input(pe_cache_.get_or_create(image_pe_key, [&]() {
+                auto 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),
@@ -671,8 +675,8 @@ namespace Anima {
                                                                               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());
+                return sd::Tensor<float>({2, 2, config.head_dim / 2, image_pos_len}, std::move(image_pe_vec));
+            }));
 
             ggml_tensor* adapter_q_pe = nullptr;
             ggml_tensor* adapter_k_pe = nullptr;
@@ -680,16 +684,16 @@ namespace Anima {
                 int64_t target_len = t5_ids->ne[0];
                 int64_t source_len = context->ne[1];
 
-                adapter_q_pe_vec = gen_1d_rope_pe_vec(target_len, 64, 10000.f);
-                adapter_k_pe_vec = gen_1d_rope_pe_vec(source_len, 64, 10000.f);
-
-                int64_t target_pos_len = static_cast<int64_t>(adapter_q_pe_vec.size()) / (2 * 2 * 32);
-                int64_t source_pos_len = static_cast<int64_t>(adapter_k_pe_vec.size()) / (2 * 2 * 32);
-
-                adapter_q_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 32, target_pos_len);
-                adapter_k_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 32, source_pos_len);
-                set_backend_tensor_data(adapter_q_pe, adapter_q_pe_vec.data());
-                set_backend_tensor_data(adapter_k_pe, adapter_k_pe_vec.data());
+                sd::PECacheKey q_key("anima.adapter_q_pe");
+                q_key.add(target_len);
+                adapter_q_pe = make_input(pe_cache_.get_or_create(q_key, [&]() {
+                    return sd::Tensor<float>({2, 2, 32, target_len}, gen_1d_rope_pe_vec(target_len, 64, 10000.f));
+                }));
+                sd::PECacheKey k_key("anima.adapter_k_pe");
+                k_key.add(source_len);
+                adapter_k_pe = make_input(pe_cache_.get_or_create(k_key, [&]() {
+                    return sd::Tensor<float>({2, 2, 32, source_len}, gen_1d_rope_pe_vec(source_len, 64, 10000.f));
+                }));
             }
 
             auto runner_ctx = get_context();

src/model/diffusion/boogu.hpp

@@ -6,6 +6,7 @@
 #include <cmath>
 #include <tuple>
 #include <vector>
+#include "core/pe_cache.h"
 
 #include "core/ggml_extend.h"
 #include "core/ggml_runner.h"
@@ -759,9 +760,9 @@ namespace Boogu {
     }
 
     struct BooguImageRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
         BooguConfig config;
         BooguImageModel boogu;
-        std::vector<float> pe_vec;
 
         BooguImageRunner(ggml_backend_t backend,
                          const String2TensorStorage& tensor_storage_map      = {},
@@ -799,7 +800,11 @@ namespace Boogu {
                 ref_latents.push_back(make_input(ref_latent_tensor));
             }
 
-            pe_vec      = finish_rope_pe(gen_boogu_pe(static_cast<int>(x->ne[1]),
+            sd::PECacheKey pe_key("boogu.pe");
+            pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+            pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size).add_shapes(ref_latents);
+            auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+                auto 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]),
@@ -808,8 +813,8 @@ namespace Boogu {
                                                           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());
+                return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+            }));
 
             auto runner_ctx  = get_context();
             ggml_tensor* out = boogu.forward(&runner_ctx, x, timesteps, context, pe, ref_latents);

src/model/diffusion/ernie_image.hpp

@@ -4,6 +4,7 @@
 #include <cinttypes>
 #include <memory>
 #include <vector>
+#include "core/pe_cache.h"
 
 #include "model/common/rope.hpp"
 #include "model/diffusion/dit.hpp"
@@ -383,9 +384,9 @@ namespace ErnieImage {
     };
 
     struct ErnieImageRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
         ErnieImageConfig config;
         ErnieImageModel ernie_image;
-        std::vector<float> pe_vec;
 
         ErnieImageRunner(ggml_backend_t backend,
                          const String2TensorStorage& tensor_storage_map      = {},
@@ -415,7 +416,11 @@ namespace ErnieImage {
             GGML_ASSERT(!context_tensor.empty());
             ggml_tensor* context = make_input(context_tensor);
 
-            pe_vec      = finish_rope_pe(Rope::gen_ernie_image_pe(static_cast<int>(x->ne[1]),
+            sd::PECacheKey pe_key("ernie_image.pe");
+            pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+            pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size);
+            auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+                auto 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]),
@@ -423,8 +428,8 @@ namespace ErnieImage {
                                                                       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());
+                return sd::Tensor<float>({config.axes_dim_sum, 1, pos_len, 2}, std::move(pe_vec));
+            }));
 
             auto runner_ctx  = get_context();
             ggml_tensor* out = ernie_image.forward(&runner_ctx, x, timesteps, context, pe);

src/model/diffusion/fibo2.hpp

@@ -1,6 +1,7 @@
 #ifndef __SD_MODEL_DIFFUSION_FIBO2_HPP__
 #define __SD_MODEL_DIFFUSION_FIBO2_HPP__
 
+#include "core/pe_cache.h"
 #include "model/diffusion/z_image.hpp"
 
 // https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_bria_fibo2.py
@@ -333,9 +334,9 @@ namespace Fibo2 {
     };
 
     struct Fibo2Runner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
         Fibo2Config config;
         Fibo2Model model;
-        std::vector<float> pe_values;
 
         Fibo2Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
             : DiffusionModelRunner(backend, prefix, weight_manager),
@@ -360,6 +361,16 @@ namespace Fibo2 {
                 auto x     = make_input(*p.x);
                 GGML_ASSERT(x->ne[3] == 1 && x->ne[2] == config.channels);
                 int count = config.gist_count(static_cast<int>(p.context->shape()[1]));
+                std::vector<ggml_tensor*> refs;
+                if (p.ref_latents && p.ref_image_params.pass_to_dit) {
+                    for (const auto& ref : *p.ref_latents) {
+                        refs.push_back(make_input(ref));
+                    }
+                }
+                sd::PECacheKey pe_key("fibo2.pe");
+                pe_key.add_shape(x).add(count).add(config.max_gist).add_shapes(refs).add(circular_x_enabled);
+                pe_key.add(circular_y_enabled);
+                auto pe  = make_input(pe_cache_.get_or_create(pe_key, [&]() {
                     Rope::Embedding embedding;
                     for (int i = 0; i < count; ++i) {
                         embedding.ids.push_back({static_cast<float>(i + 1), 0.f, 0.f});
@@ -371,17 +382,14 @@ namespace Fibo2 {
                         embedding.ids.insert(embedding.ids.end(), ids.begin(), ids.end());
                     };
                     append_image(x, static_cast<int>(config.max_gist + 1));
-                std::vector<ggml_tensor*> refs;
-                if (p.ref_latents && p.ref_image_params.pass_to_dit) {
-                    for (const auto& ref : *p.ref_latents) {
-                        refs.push_back(make_input(ref));
-                        append_image(refs.back(), static_cast<int>(config.max_gist + 1 + refs.size()));
-                    }
+                    for (size_t i = 0; i < refs.size(); ++i) {
+                        append_image(refs[i], static_cast<int>(config.max_gist + 2 + i));
                     }
                     embedding.values     = Rope::embed_nd(embedding.ids, 1, 256.f, {32, 48, 48}, embedding.layout, &embedding.frequencies);
-                pe_values        = finish_rope_pe(std::move(embedding));
-                auto pe          = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 64, static_cast<int64_t>(pe_values.size() / 256));
-                set_backend_tensor_data(pe, pe_values.data());
+                    auto values          = finish_rope_pe(std::move(embedding));
+                    const int64_t tokens = static_cast<int64_t>(values.size() / 256);
+                    return sd::Tensor<float>({2, 2, 64, tokens}, std::move(values));
+                }));
                 auto ctx = get_context();
                 auto out = model.forward(&ctx, x, make_input(*p.timesteps), make_input(*p.context), pe, count, refs);
                 ggml_build_forward_expand(graph, out);

src/model/diffusion/flux.hpp

@@ -6,6 +6,7 @@
 #include <vector>
 #include "core/ggml_extend_backend.h"
 #include "core/ggml_tensor_utils.h"
+#include "core/pe_cache.h"
 
 #include "core/util.h"
 #include "model/adapter/pulid.hpp"
@@ -1398,10 +1399,11 @@ namespace Flux {
     };
 
     struct FluxRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
+
     public:
         FluxConfig config;
         Flux flux;
-        std::vector<float> pe_vec;
         std::vector<float> mod_index_arange_vec;
         std::vector<float> dct_vec;
         sd::Tensor<float> guidance_tensor;
@@ -1548,7 +1550,12 @@ namespace Flux {
             } else if (version == VERSION_OVIS_IMAGE) {
                 txt_arange_dims = {1, 2};
             }
-            pe_vec      = finish_rope_pe(Rope::gen_flux_pe(static_cast<int>(x->ne[1]),
+            sd::PECacheKey pe_key("flux.pe");
+            pe_key.add(version).add_shape(x).add_shape(context).add(config.patch_size).add(txt_arange_dims);
+            pe_key.add_shapes(ref_latents).add(ref_index_mode).add(config.ref_index_scale).add(config.theta).add(config.axes_dim);
+            pe_key.add(circular_x_enabled).add(circular_y_enabled);
+            auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+                auto 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]),
@@ -1560,13 +1567,9 @@ namespace Flux {
                                                                    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);
-            // pe->data = pe_vec.data();
-            // print_ggml_tensor(pe);
-            // pe->data = nullptr;
-            set_backend_tensor_data(pe, pe_vec.data());
+                int64_t pos_len = static_cast<int64_t>(pe_vec.size() / config.axes_dim_sum / 2);
+                return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+            }));
 
             if (version == VERSION_CHROMA_RADIANCE) {
                 int patch_size     = config.patch_size;

src/model/diffusion/flux3.hpp

@@ -1,6 +1,7 @@
 #ifndef __SD_MODEL_DIFFUSION_FLUX3_HPP__
 #define __SD_MODEL_DIFFUSION_FLUX3_HPP__
 
+#include "core/pe_cache.h"
 #include "model/diffusion/flux.hpp"
 
 namespace Flux3 {
@@ -183,9 +184,9 @@ namespace Flux3 {
     };
 
     struct Flux3Runner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
         Flux3Config config;
         Model model;
-        std::vector<float> pe_values;
         sd::Tensor<float> zero = sd::Tensor<float>::zeros({1});
 
         Flux3Runner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
@@ -213,6 +214,19 @@ namespace Flux3 {
                 auto x    = make_input(*params.x);
                 auto text = make_input(*params.context);
                 GGML_ASSERT(x->ne[3] == 1 && text->ne[2] == 1);
+                std::vector<ggml_tensor*> refs;
+                ggml_tensor* references = nullptr;
+                if (params.ref_latents && params.ref_image_params.pass_to_dit) {
+                    for (const auto& ref : *params.ref_latents) {
+                        auto r = make_input(ref);
+                        refs.push_back(r);
+                        r          = DiT::patchify(compute_ctx, r, 1, 1);
+                        references = references ? ggml_concat(compute_ctx, references, r, 1) : r;
+                    }
+                }
+                sd::PECacheKey pe_key("flux3.pe");
+                pe_key.add_shape(x).add_shape(text).add_shapes(refs).add(circular_x_enabled).add(circular_y_enabled);
+                auto pe         = make_input(pe_cache_.get_or_create(pe_key, [&]() {
                     Rope::Embedding embedding;
                     auto append_image = [&](int64_t w, int64_t h, float index) {
                         embedding.positions.append_image(static_cast<int>(h), static_cast<int>(w), 1, 0.5f, 0.5f);
@@ -223,24 +237,18 @@ namespace Flux3 {
                         }
                     };
                     append_image(x->ne[0], x->ne[1], 0.f);
-                ggml_tensor* references = nullptr;
-                if (params.ref_latents && params.ref_image_params.pass_to_dit) {
-                    int index = 0;
-                    for (const auto& ref : *params.ref_latents) {
-                        auto r = make_input(ref);
-                        append_image(r->ne[0], r->ne[1], -float(++index));
-                        r          = DiT::patchify(compute_ctx, r, 1, 1);
-                        references = references ? ggml_concat(compute_ctx, references, r, 1) : r;
-                    }
+                    for (size_t i = 0; i < refs.size(); ++i) {
+                        append_image(refs[i]->ne[0], refs[i]->ne[1], -float(i + 1));
                     }
                     for (int64_t i = 0; i < text->ne[1]; ++i) {
                         embedding.ids.push_back({0.f, 0.f, 0.f, static_cast<float>(i)});
                     }
                     embedding.positions.append_tokens(text->ne[1]);
                     embedding.values     = Rope::embed_nd(embedding.ids, 1, 10000.f, {32, 32, 32, 32}, embedding.layout, &embedding.frequencies);
-                pe_values        = finish_rope_pe(std::move(embedding));
-                auto pe          = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 64, pe_values.size() / 256);
-                set_backend_tensor_data(pe, pe_values.data());
+                    auto values          = finish_rope_pe(std::move(embedding));
+                    const int64_t tokens = static_cast<int64_t>(values.size() / 256);
+                    return sd::Tensor<float>({2, 2, 64, tokens}, std::move(values));
+                }));
                 auto runner_ctx = get_context();
                 auto out        = model.forward(&runner_ctx, x, text, references, make_input(*params.timesteps), make_input(zero),
                                                 config.flow_map ? make_input(*extra->target_timestep) : nullptr, pe);

src/model/diffusion/hunyuan.hpp

@@ -3,6 +3,7 @@
 
 #include <cinttypes>
 #include <memory>
+#include "core/pe_cache.h"
 
 #include "model/common/block.hpp"
 #include "model/diffusion/flux.hpp"
@@ -540,10 +541,11 @@ namespace Hunyuan {
     };
 
     struct HunyuanVideoRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
+
     public:
         HunyuanVideoConfig config;
         HunyuanVideoModel hunyuan_video;
-        std::vector<float> pe_vec;
         SDVersion version;
 
         HunyuanVideoRunner(ggml_backend_t backend,
@@ -605,7 +607,12 @@ namespace Hunyuan {
             if (vision != nullptr) {
                 text_len += static_cast<int>(vision->ne[1]);
             }
-            pe_vec          = Rope::gen_hunyuan_video_pe(static_cast<int>(x->ne[2]),
+            sd::PECacheKey pe_key("hunyuan.pe");
+            pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+            pe_key.add(circular_x_enabled).add(circular_y_enabled).add(std::get<0>(config.patch_size));
+            pe_key.add(std::get<1>(config.patch_size)).add(std::get<2>(config.patch_size)).add(text_len);
+            auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+                auto pe_vec     = Rope::gen_hunyuan_video_pe(static_cast<int>(x->ne[2]),
                                                              static_cast<int>(x->ne[1]),
                                                              static_cast<int>(x->ne[0]),
                                                              std::get<0>(config.patch_size),
@@ -616,12 +623,8 @@ namespace Hunyuan {
                                                              config.theta,
                                                              config.axes_dim);
                 int64_t pos_len = static_cast<int64_t>(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);
-            // pe->data = pe_vec.data();
-            // print_ggml_tensor(pe, true, "pe");
-            // pe->data = nullptr;
-            set_backend_tensor_data(pe, pe_vec.data());
+                return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+            }));
 
             auto runner_ctx = get_context();
 

src/model/diffusion/ideogram4.hpp

@@ -8,6 +8,7 @@
 #include <memory>
 #include <string>
 #include <vector>
+#include "core/pe_cache.h"
 
 #include "core/ggml_extend.h"
 #include "core/ggml_graph_cut.h"
@@ -422,6 +423,8 @@ namespace Ideogram4 {
     };
 
     class Ideogram4Runner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
+
     protected:
         bool should_use_uncond_model(const DiffusionParams& diffusion_params) const {
             return has_uncond_model &&
@@ -436,7 +439,6 @@ namespace Ideogram4 {
         Ideogram4Transformer uncond_model;
         bool has_uncond_model = false;
         std::string uncond_prefix;
-        std::vector<float> pe_vec;
         std::vector<int32_t> image_indicator_vec;
 
         Ideogram4Runner(ggml_backend_t backend,
@@ -496,15 +498,19 @@ namespace Ideogram4 {
             int64_t head_dim = config.emb_dim / config.num_heads;
 
             auto runner_ctx = get_context();
-            pe_vec          = finish_rope_pe(gen_ideogram4_pe(static_cast<int>(grid_h),
+            sd::PECacheKey pe_key("ideogram4.pe");
+            pe_key.add(grid_h).add(grid_w).add(x->ne[3]).add(context_len).add(head_dim).add(pos_len).add(config.rope_theta);
+            pe_key.add(config.mrope_section).add(circular_x_enabled).add(circular_y_enabled);
+            auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+                auto 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());
+                return sd::Tensor<float>({2, 2, head_dim / 2, pos_len}, std::move(pe_vec));
+            }));
 
             image_indicator_vec.assign(static_cast<size_t>(pos_len), 1);
             for (int64_t i = 0; i < context_len; ++i) {

src/model/diffusion/iris.hpp

@@ -3,6 +3,7 @@
 
 #include <array>
 #include <stdexcept>
+#include "core/pe_cache.h"
 #include "model/diffusion/pid.hpp"
 
 namespace Iris {
@@ -382,11 +383,9 @@ namespace Iris {
     }
 
     struct IrisRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
         IrisConfig config;
         IrisModel model;
-        std::vector<float> joint_pe_data;
-        std::vector<float> pixel_pe_data;
-        std::vector<float> pixel_pos_data;
         std::vector<float> text_mask_data;
 
         IrisRunner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
@@ -418,24 +417,35 @@ namespace Iris {
                 GGML_ASSERT(h % config.patch_size == 0 && w % config.patch_size == 0);
                 GGML_ASSERT(text->ne[0] == config.text_dim * config.text_layers && text->ne[1] == nt);
                 GGML_ASSERT(inputs.y->numel() == nt && x->ne[3] == 1);
-                joint_pe_data = Pid::make_rope_1d(nt, static_cast<int>(config.head_dim), 10000.f);
+                sd::PECacheKey joint_key("iris.joint_pe");
+                joint_key.add(h).add(w).add(nt).add(config.patch_size).add(config.head_dim);
+                auto pe = make_input(pe_cache_.get_or_create(joint_key, [&]() {
+                    auto values = Pid::make_rope_1d(nt, static_cast<int>(config.head_dim), 10000.f);
                     auto img    = image_rope(h / config.patch_size, w / config.patch_size, config.head_dim);
-                joint_pe_data.insert(joint_pe_data.end(), img.begin(), img.end());
-                pixel_pe_data  = image_rope(h / config.patch_size, w / config.patch_size, config.pixel_attn_dim / config.pixel_heads);
-                pixel_pos_data = Pid::make_pixel_abs_pos(h, w, static_cast<int>(config.pixel_dim));
+                    values.insert(values.end(), img.begin(), img.end());
+                    const int64_t tokens = nt + static_cast<int64_t>(img.size() / (2 * config.head_dim));
+                    return sd::Tensor<float>({2, 2, config.head_dim / 2, tokens}, std::move(values));
+                }));
+                sd::PECacheKey pixel_key("iris.pixel_pe");
+                pixel_key.add(h).add(w).add(config.patch_size).add(config.pixel_attn_dim).add(config.pixel_heads);
+                auto pp = make_input(pe_cache_.get_or_create(pixel_key, [&]() {
+                    return sd::Tensor<float>({2, 2, config.pixel_attn_dim / config.pixel_heads / 2,
+                                              (h / config.patch_size) * (w / config.patch_size)},
+                                             image_rope(h / config.patch_size, w / config.patch_size, config.pixel_attn_dim / config.pixel_heads));
+                }));
+                sd::PECacheKey pos_key("iris.pixel_pos");
+                pos_key.add(h).add(w).add(config.pixel_dim);
+                auto pos = make_input(pe_cache_.get_or_create(pos_key, [&]() {
+                    return sd::Tensor<float>({config.pixel_dim, h * w, 1},
+                                             Pid::make_pixel_abs_pos(h, w, static_cast<int>(config.pixel_dim)));
+                }));
                 text_mask_data.resize(nt * nt);
                 for (int q = 0; q < nt; ++q) {
                     for (int k = 0; k < nt; ++k) {
                         text_mask_data[q * nt + k] = inputs.y->values()[k] != 0.f || q == k ? 0.f : -INFINITY;
                     }
                 }
-                auto pe   = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, nt + img.size() / (2 * config.head_dim));
-                auto pp   = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.pixel_attn_dim / config.pixel_heads / 2, (h / config.patch_size) * (w / config.patch_size));
-                auto pos  = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.pixel_dim, h * w, 1);
                 auto mask = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, nt, nt);
-                set_backend_tensor_data(pe, joint_pe_data.data());
-                set_backend_tensor_data(pp, pixel_pe_data.data());
-                set_backend_tensor_data(pos, pixel_pos_data.data());
                 set_backend_tensor_data(mask, text_mask_data.data());
                 auto ctx = get_context();
                 auto out = model.forward(&ctx, x, t, text, mask, pe, pp, pos);

src/model/diffusion/krea2.hpp

@@ -11,6 +11,7 @@
 #include <string>
 #include <utility>
 #include <vector>
+#include "core/pe_cache.h"
 
 #include "core/ggml_extend.h"
 #include "core/ggml_graph_cut.h"
@@ -714,9 +715,9 @@ namespace Krea2 {
     }
 
     struct Krea2Runner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
         Krea2Config config;
         Krea2Model model;
-        std::vector<float> pe_vec;
 
         Krea2Runner(ggml_backend_t backend,
                     const String2TensorStorage& tensor_storage_map      = {},
@@ -754,7 +755,12 @@ namespace Krea2 {
                 ref_latents.push_back(make_input(ref_latent_tensor));
             }
 
-            pe_vec      = finish_rope_pe(gen_krea2_pe(static_cast<int>(x->ne[1]),
+            sd::PECacheKey pe_key("krea2.pe");
+            pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+            pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size).add_shapes(ref_latents);
+            pe_key.add(ref_image_params.ref_index_mode);
+            auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+                auto 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]),
@@ -764,8 +770,8 @@ namespace Krea2 {
                                                           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());
+                return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+            }));
 
             auto runner_ctx  = get_context();
             ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents, ref_image_params.force_ref_timestep_zero);

src/model/diffusion/lens.hpp

@@ -4,6 +4,7 @@
 #include <cinttypes>
 #include <memory>
 #include <vector>
+#include "core/pe_cache.h"
 
 #include "model/common/block.hpp"
 #include "model/common/rope.hpp"
@@ -352,9 +353,9 @@ namespace Lens {
     };
 
     struct LensRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
         LensConfig config;
         LensModel lens;
-        std::vector<float> pe_vec;
 
         LensRunner(ggml_backend_t backend,
                    const String2TensorStorage& tensor_storage_map      = {},
@@ -384,15 +385,19 @@ namespace Lens {
             GGML_ASSERT(!context_tensor.empty());
             ggml_tensor* context = make_input(context_tensor);
 
-            pe_vec      = finish_rope_pe(Rope::gen_lens_pe(static_cast<int>(x->ne[1]),
+            sd::PECacheKey pe_key("lens.pe");
+            pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+            pe_key.add(circular_x_enabled).add(circular_y_enabled);
+            auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+                auto 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());
+                return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+            }));
 
             auto runner_ctx  = get_context();
             ggml_tensor* out = lens.forward(&runner_ctx, x, timesteps, context, pe);

src/model/diffusion/lingbot_video.hpp

@@ -9,6 +9,7 @@
 #include <set>
 #include <string>
 #include <vector>
+#include "core/pe_cache.h"
 
 #include "core/util.h"
 #include "model/common/rope.hpp"
@@ -606,9 +607,9 @@ namespace LingBotVideo {
     };
 
     struct LingBotVideoRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
         LingBotVideoConfig config;
         LingBotVideoModel lingbot_video;
-        std::vector<float> pe_vec;
 
         LingBotVideoRunner(ggml_backend_t backend,
                            const String2TensorStorage& tensor_storage_map      = {},
@@ -647,7 +648,11 @@ namespace LingBotVideo {
             int64_t T = x->ne[2];
             int64_t N = x_tensor.shape()[4];
             GGML_ASSERT(N == 1);
-            pe_vec      = Rope::gen_lingbot_video_pe(static_cast<int>(T),
+            sd::PECacheKey pe_key("lingbot_video.pe");
+            pe_key.add(T).add(H).add(W).add(N).add_shape(context).add(config.patch_t).add(config.patch_h).add(config.patch_w);
+            pe_key.add(config.rope_theta).add(config.axes_dim).add(config.axes_dim_sum);
+            auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+                auto pe_vec = Rope::gen_lingbot_video_pe(static_cast<int>(T),
                                                          static_cast<int>(H),
                                                          static_cast<int>(W),
                                                          config.patch_t,
@@ -658,8 +663,8 @@ namespace LingBotVideo {
                                                          config.rope_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());
+                return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+            }));
 
             auto runner_ctx  = get_context();
             ggml_tensor* out = lingbot_video.forward(&runner_ctx, x, timesteps, context, pe);

src/model/diffusion/llada_image.hpp

@@ -3,6 +3,7 @@
 
 #include <algorithm>
 #include <cinttypes>
+#include "core/pe_cache.h"
 
 #include "core/ggml_extend.h"
 #include "core/ggml_runner.h"
@@ -379,10 +380,11 @@ namespace LLaDAImage {
     };
 
     struct LLaDAImageRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
+
     public:
         LLaDAImageConfig config;
         LLaDAImageModel llada_image;
-        std::vector<float> pe_vec;
 
         LLaDAImageRunner(ggml_backend_t backend,
                          const String2TensorStorage& tensor_storage_map      = {},
@@ -412,7 +414,11 @@ namespace LLaDAImage {
             GGML_ASSERT(!context_tensor.empty());
             ggml_tensor* context = make_input(context_tensor);
 
-            pe_vec      = finish_rope_pe(Rope::gen_llada_image_pe(static_cast<int>(x->ne[1]),
+            sd::PECacheKey pe_key("llada_image.pe");
+            pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+            pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size);
+            auto pe         = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+                auto 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]),
@@ -421,8 +427,8 @@ namespace LLaDAImage {
                                                                       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());
+                return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+            }));
             auto runner_ctx = get_context();
 
             ggml_tensor* out = llada_image.forward(&runner_ctx, x, timesteps, context, pe);
@@ -459,7 +465,11 @@ namespace LLaDAImage {
             ggml_tensor* source    = make_input(source_tensor);
             GGML_ASSERT(x->ne[3] == 1);
 
-            pe_vec      = finish_rope_pe(Rope::gen_llada_image_edit_pe(static_cast<int>(x->ne[1]),
+            sd::PECacheKey pe_key("llada_image.pe1");
+            pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+            pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size).add_shape(semantic);
+            auto pe         = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+                auto 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]),
@@ -468,8 +478,8 @@ namespace LLaDAImage {
                                                                            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());
+                return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+            }));
             auto runner_ctx = get_context();
 
             int64_t W   = x->ne[0];

src/model/diffusion/ltxv.hpp

@@ -11,6 +11,7 @@
 #include <vector>
 #include "core/ggml_extend_backend.h"
 #include "core/ggml_tensor_utils.h"
+#include "core/pe_cache.h"
 
 #include "model/common/block.hpp"
 #include "model/common/rope.hpp"
@@ -1736,14 +1737,9 @@ namespace LTXV {
     };
 
     struct LTXAVRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
         LTXAVConfig config;
         LTXAVModelBlock model;
-        std::vector<float> video_pe_vec;
-        std::vector<float> audio_pe_vec;
-        std::vector<float> video_cross_pe_vec;
-        std::vector<float> audio_cross_pe_vec;
-        std::vector<float> connector_pe_vec;
-        std::vector<float> audio_connector_pe_vec;
         sd::Tensor<float> vx_input_cache;
         sd::Tensor<float> ax_input_cache;
 
@@ -1853,6 +1849,14 @@ namespace LTXV {
             float video_frame_rate    = frame_rate > 0.f ? frame_rate : 24.f;
             int64_t video_token_count = vx->ne[0] * vx->ne[1] * vx->ne[2];
             bool has_video_positions  = !video_positions_tensor.empty();
+            sd::PECacheKey video_pe_key("ltxav_video_pe");
+            video_pe_key.add_shape(vx).add(video_frame_rate).add(video_positions_tensor).add(config.positional_embedding_theta);
+            video_pe_key.add(config.causal_temporal_positioning).add(config.use_middle_indices_grid).add(config.hidden_size);
+            video_pe_key.add(config.num_attention_heads).add(config.attention_head_dim).add(config.positional_embedding_max_pos);
+            video_pe_key.add(std::get<0>(config.vae_scale_factors)).add(std::get<1>(config.vae_scale_factors));
+            video_pe_key.add(std::get<2>(config.vae_scale_factors));
+            auto video_pe = make_input(pe_cache_.get_or_create(video_pe_key, [&]() {
+                std::vector<float> video_pe_vec;
                 if (has_video_positions) {
                     GGML_ASSERT(video_positions_tensor.shape()[2] == video_token_count);
                     video_pe_vec = build_video_rope_matrix_from_positions(video_positions_tensor,
@@ -1874,25 +1878,36 @@ namespace LTXV {
                                                            config.causal_temporal_positioning,
                                                            config.use_middle_indices_grid);
                 }
-            auto video_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.attention_head_dim / 2, video_token_count * config.num_attention_heads);
-            ggml_set_name(video_pe, "ltxav_video_pe");
-            set_backend_tensor_data(video_pe, video_pe_vec.data());
+                return sd::Tensor<float>({2, 2, config.attention_head_dim / 2, video_token_count * config.num_attention_heads}, std::move(video_pe_vec));
+            }));
 
             ggml_tensor* audio_pe       = nullptr;
             ggml_tensor* video_cross_pe = nullptr;
             ggml_tensor* audio_cross_pe = nullptr;
             if (ax != nullptr && ggml_nelements(ax) > 0 && ax->ne[1] > 0) {
-                audio_pe_vec = build_audio_rope_matrix(ax->ne[1],
+                sd::PECacheKey audio_pe_key("ltxav_audio_pe");
+                audio_pe_key.add(ax->ne[1]).add(config.audio_hidden_size).add(config.audio_num_attention_heads);
+                audio_pe_key.add(config.audio_attention_head_dim).add(config.positional_embedding_theta);
+                audio_pe_key.add(config.audio_positional_embedding_max_pos[0]).add(config.use_middle_indices_grid);
+                audio_pe = make_input(pe_cache_.get_or_create(audio_pe_key, [&]() {
+                    auto audio_pe_vec = build_audio_rope_matrix(ax->ne[1],
                                                                 static_cast<int>(config.audio_hidden_size),
                                                                 static_cast<int>(config.audio_num_attention_heads),
                                                                 config.positional_embedding_theta,
                                                                 config.audio_positional_embedding_max_pos[0],
                                                                 config.use_middle_indices_grid);
-                audio_pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads);
-                ggml_set_name(audio_pe, "ltxav_audio_pe");
-                set_backend_tensor_data(audio_pe, audio_pe_vec.data());
+                    return sd::Tensor<float>({2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads}, std::move(audio_pe_vec));
+                }));
 
                 int temporal_max_pos = std::max(config.positional_embedding_max_pos[0], config.audio_positional_embedding_max_pos[0]);
+                sd::PECacheKey video_cross_pe_key("ltxav_video_cross_pe");
+                video_cross_pe_key.add_shape(vx).add(video_frame_rate).add(video_positions_tensor);
+                video_cross_pe_key.add(config.positional_embedding_theta).add(config.causal_temporal_positioning);
+                video_cross_pe_key.add(config.use_middle_indices_grid).add(config.audio_cross_attention_dim);
+                video_cross_pe_key.add(config.audio_num_attention_heads).add(config.audio_attention_head_dim);
+                video_cross_pe_key.add(temporal_max_pos).add(std::get<0>(config.vae_scale_factors));
+                video_cross_pe = make_input(pe_cache_.get_or_create(video_cross_pe_key, [&]() {
+                    std::vector<float> video_cross_pe_vec;
                     if (has_video_positions) {
                         video_cross_pe_vec = build_video_temporal_rope_matrix_from_positions(video_positions_tensor,
                                                                                              static_cast<int>(config.audio_cross_attention_dim),
@@ -1913,19 +1928,22 @@ namespace LTXV {
                                                                               config.causal_temporal_positioning,
                                                                               true);
                     }
-                video_cross_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, video_token_count * config.audio_num_attention_heads);
-                ggml_set_name(video_cross_pe, "ltxav_video_cross_pe");
-                set_backend_tensor_data(video_cross_pe, video_cross_pe_vec.data());
+                    return sd::Tensor<float>({2, 2, config.audio_attention_head_dim / 2, video_token_count * config.audio_num_attention_heads}, std::move(video_cross_pe_vec));
+                }));
 
-                audio_cross_pe_vec = build_audio_rope_matrix(ax->ne[1],
+                sd::PECacheKey audio_cross_pe_key("ltxav_audio_cross_pe");
+                audio_cross_pe_key.add(ax->ne[1]).add(config.audio_cross_attention_dim).add(config.audio_num_attention_heads);
+                audio_cross_pe_key.add(config.audio_attention_head_dim).add(config.positional_embedding_theta);
+                audio_cross_pe_key.add(temporal_max_pos);
+                audio_cross_pe = make_input(pe_cache_.get_or_create(audio_cross_pe_key, [&]() {
+                    auto audio_cross_pe_vec = build_audio_rope_matrix(ax->ne[1],
                                                                       static_cast<int>(config.audio_cross_attention_dim),
                                                                       static_cast<int>(config.audio_num_attention_heads),
                                                                       config.positional_embedding_theta,
                                                                       temporal_max_pos,
                                                                       true);
-                audio_cross_pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads);
-                ggml_set_name(audio_cross_pe, "ltxav_audio_cross_pe");
-                set_backend_tensor_data(audio_cross_pe, audio_cross_pe_vec.data());
+                    return sd::Tensor<float>({2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads}, std::move(audio_cross_pe_vec));
+                }));
             }
 
             bool needs_video_connector_pe =
@@ -1941,10 +1959,13 @@ namespace LTXV {
                 int64_t target_len   = std::max<int64_t>(1024, seq_len);
                 int64_t duplications = (target_len + config.connector_num_registers - 1) / config.connector_num_registers;
                 int64_t full_len     = seq_len + duplications * config.connector_num_registers - seq_len;
-                connector_pe_vec     = build_1d_rope_matrix(full_len, static_cast<int>(config.connector_hidden_size), static_cast<int>(config.connector_num_heads), 10000.f, 4096.f, true);
-                video_connector_pe   = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.connector_head_dim / 2, full_len * config.connector_num_heads);
-                ggml_set_name(video_connector_pe, "ltxav_video_connector_pe");
-                set_backend_tensor_data(video_connector_pe, connector_pe_vec.data());
+                sd::PECacheKey video_connector_pe_key("ltxav_video_connector_pe");
+                video_connector_pe_key.add(full_len).add(config.connector_hidden_size).add(config.connector_num_heads);
+                video_connector_pe_key.add(config.connector_head_dim);
+                video_connector_pe = make_input(pe_cache_.get_or_create(video_connector_pe_key, [&]() {
+                    auto connector_pe_vec = build_1d_rope_matrix(full_len, static_cast<int>(config.connector_hidden_size), static_cast<int>(config.connector_num_heads), 10000.f, 4096.f, true);
+                    return sd::Tensor<float>({2, 2, config.connector_head_dim / 2, full_len * config.connector_num_heads}, std::move(connector_pe_vec));
+                }));
             }
 
             bool run_audio_context =
@@ -1965,10 +1986,13 @@ namespace LTXV {
                 int64_t target_len   = std::max<int64_t>(1024, seq_len);
                 int64_t duplications = (target_len + config.audio_connector_num_registers - 1) / config.audio_connector_num_registers;
                 int64_t full_len     = seq_len + duplications * config.audio_connector_num_registers - seq_len;
-                audio_connector_pe_vec = build_1d_rope_matrix(full_len, static_cast<int>(config.audio_connector_hidden_size), static_cast<int>(config.audio_connector_num_heads), 10000.f, 4096.f, true);
-                audio_connector_pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_connector_head_dim / 2, full_len * config.audio_connector_num_heads);
-                ggml_set_name(audio_connector_pe, "ltxav_audio_connector_pe");
-                set_backend_tensor_data(audio_connector_pe, audio_connector_pe_vec.data());
+                sd::PECacheKey audio_connector_pe_key("ltxav_audio_connector_pe");
+                audio_connector_pe_key.add(full_len).add(config.audio_connector_hidden_size).add(config.audio_connector_num_heads);
+                audio_connector_pe_key.add(config.audio_connector_head_dim);
+                audio_connector_pe = make_input(pe_cache_.get_or_create(audio_connector_pe_key, [&]() {
+                    auto audio_connector_pe_vec = build_1d_rope_matrix(full_len, static_cast<int>(config.audio_connector_hidden_size), static_cast<int>(config.audio_connector_num_heads), 10000.f, 4096.f, true);
+                    return sd::Tensor<float>({2, 2, config.audio_connector_head_dim / 2, full_len * config.audio_connector_num_heads}, std::move(audio_connector_pe_vec));
+                }));
             }
 
             auto runner_ctx = get_context();

src/model/diffusion/mage_flow.hpp

@@ -3,6 +3,7 @@
 
 #include <cmath>
 #include <memory>
+#include "core/pe_cache.h"
 
 #include "model/diffusion/qwen_image.hpp"
 
@@ -57,11 +58,12 @@ namespace MageFlow {
     };
 
     struct MageFlowRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
+
     public:
         Qwen::QwenImageConfig config;
         Qwen::QwenImageModel mage_flow;
         std::shared_ptr<MageFlowTimestepProjEmbeddings> time_text_embed;
-        std::vector<float> pe_vec;
 
         MageFlowRunner(ggml_backend_t backend,
                        const String2TensorStorage& tensor_storage_map      = {},
@@ -110,7 +112,11 @@ namespace MageFlow {
             }
 
             int batch_size = static_cast<int>(x->ne[3]);
-            pe_vec         = finish_rope_pe(Rope::gen_mage_flow_pe(static_cast<int>(x->ne[1]),
+            sd::PECacheKey pe_key("mage_flow.pe");
+            pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+            pe_key.add(circular_x_enabled).add(circular_y_enabled).add(batch_size).add_shapes(ref_latents);
+            auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+                auto 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]),
@@ -118,8 +124,8 @@ namespace MageFlow {
                                                                     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());
+                return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+            }));
 
             time_text_embed->prepare(timesteps_tensor);
             auto runner_ctx = get_context();

src/model/diffusion/ming_image.hpp

@@ -1,6 +1,7 @@
 #ifndef __SD_MODEL_DIFFUSION_MING_IMAGE_HPP__
 #define __SD_MODEL_DIFFUSION_MING_IMAGE_HPP__
 
+#include "core/pe_cache.h"
 #include "z_image.hpp"
 
 namespace MingImage {
@@ -72,9 +73,9 @@ namespace MingImage {
     };
 
     struct MingImageRunner : DiffusionModelRunner {
+        sd::PECache pe_cache_;
         MingImageConfig config;
         MingImageModel model;
-        std::vector<float> pe_values;
 
         MingImageRunner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
             : DiffusionModelRunner(backend, prefix, weight_manager),
@@ -109,6 +110,10 @@ namespace MingImage {
                 const int64_t n_txt = context->ne[1] + direct->ne[1];
                 const int64_t n_img = ((x->ne[0] + config.patch_size - 1) / config.patch_size) *
                                       ((x->ne[1] + config.patch_size - 1) / config.patch_size);
+                sd::PECacheKey pe_key("ming_image.pe");
+                pe_key.add_shape(x).add(n_txt).add(config.patch_size).add(config.theta).add(config.axes_dim);
+                pe_key.add(config.axes_dim_sum).add(circular_x_enabled).add(circular_y_enabled);
+                auto pe  = make_input(pe_cache_.get_or_create(pe_key, [&]() {
                     auto padded = finish_rope_pe(Rope::gen_z_image_pe(
                         static_cast<int>(x->ne[1]), static_cast<int>(x->ne[0]), config.patch_size, 1,
                         static_cast<int>(n_txt), ZImage::SEQ_MULTI_OF, {}, Rope::RefIndexMode::FIXED,
@@ -117,11 +122,11 @@ namespace MingImage {
                     // retaining the padded caption length used to position image tokens.
                     const size_t stride      = config.axes_dim_sum * 2;
                     const int64_t padded_txt = n_txt + Rope::bound_mod(static_cast<int>(n_txt), ZImage::SEQ_MULTI_OF);
-                pe_values.assign(padded.begin(), padded.begin() + n_txt * stride);
+                    std::vector<float> pe_values(padded.begin(), padded.begin() + n_txt * stride);
                     pe_values.insert(pe_values.end(), padded.begin() + padded_txt * stride,
                                      padded.begin() + (padded_txt + n_img) * stride);
-                auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, n_txt + n_img);
-                set_backend_tensor_data(pe, pe_values.data());
+                    return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, n_txt + n_img}, std::move(pe_values));
+                }));
                 auto ctx = get_context();
                 auto out = model.forward(&ctx, x, t, context, direct, pe);
                 ggml_build_forward_expand(gf, out);

src/model/diffusion/minit2i.hpp

@@ -9,6 +9,7 @@
 #include <memory>
 #include <string>
 #include <vector>
+#include "core/pe_cache.h"
 
 #include "core/ggml_extend.h"
 #include "core/ggml_runner.h"
@@ -473,20 +474,9 @@ namespace MiniT2I {
     };
 
     struct MiniT2IRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
         MiniT2IConfig config;
         MMJiT model;
-        ggml_context* position_cache_ctx            = nullptr;
-        ggml_backend_buffer_t position_cache_buffer = nullptr;
-        ggml_tensor* cached_pos_embed               = nullptr;
-        ggml_tensor* cached_txt_pe                  = nullptr;
-        ggml_tensor* cached_joint_pe                = nullptr;
-        int64_t cached_img_side                     = -1;
-        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      = {},
                       const std::string prefix                            = "",
@@ -497,10 +487,6 @@ namespace MiniT2I {
             model.init(params_ctx, tensor_storage_map, this->prefix);
         }
 
-        ~MiniT2IRunner() override {
-            free_position_cache();
-        }
-
         std::string get_desc() override {
             return "MiniT2I";
         }
@@ -509,75 +495,6 @@ namespace MiniT2I {
             model.get_param_tensors(tensors, prefix);
         }
 
-        void free_position_cache() {
-            if (position_cache_buffer != nullptr) {
-                ggml_backend_buffer_free(position_cache_buffer);
-                position_cache_buffer = nullptr;
-            }
-            if (position_cache_ctx != nullptr) {
-                ggml_free(position_cache_ctx);
-                position_cache_ctx = nullptr;
-            }
-            cached_pos_embed   = nullptr;
-            cached_txt_pe      = nullptr;
-            cached_joint_pe    = nullptr;
-            cached_img_side    = -1;
-            cached_txt_len     = -1;
-            cached_hidden_size = -1;
-            cached_head_dim    = -1;
-        }
-
-        void ensure_position_cache(int64_t img_side, int64_t txt_len) {
-            if (cached_img_side == img_side &&
-                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) {
-                return;
-            }
-
-            free_position_cache();
-
-            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    = 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());
-
-            ggml_init_params params;
-            params.mem_size    = static_cast<size_t>(3 * ggml_tensor_overhead());
-            params.mem_buffer  = nullptr;
-            params.no_alloc    = true;
-            position_cache_ctx = ggml_init(params);
-            GGML_ASSERT(position_cache_ctx != nullptr);
-
-            cached_pos_embed = ggml_new_tensor_3d(position_cache_ctx, GGML_TYPE_F32, config.hidden_size, img_side * img_side, 1);
-            ggml_set_name(cached_pos_embed, "minit2i.pos_embed");
-            cached_txt_pe = ggml_new_tensor_4d(position_cache_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, txt_len);
-            ggml_set_name(cached_txt_pe, "minit2i.txt_pe");
-            cached_joint_pe = ggml_new_tensor_4d(position_cache_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, txt_len + img_side * img_side);
-            ggml_set_name(cached_joint_pe, "minit2i.joint_pe");
-
-            position_cache_buffer = ggml_backend_alloc_ctx_tensors(position_cache_ctx, runtime_backend);
-            GGML_ASSERT(position_cache_buffer != nullptr);
-            ggml_backend_buffer_set_usage(position_cache_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
-            ggml_backend_tensor_set(cached_pos_embed, pos_embed_vec.data(), 0, ggml_nbytes(cached_pos_embed));
-            ggml_backend_tensor_set(cached_txt_pe, txt_pe_vec.data(), 0, ggml_nbytes(cached_txt_pe));
-            ggml_backend_tensor_set(cached_joint_pe, joint_pe_vec.data(), 0, ggml_nbytes(cached_joint_pe));
-            ggml_backend_synchronize(runtime_backend);
-
-            cached_img_side    = img_side;
-            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,
                                  const sd::Tensor<float>& timesteps_tensor,
                                  const sd::Tensor<float>& context_tensor,
@@ -592,10 +509,29 @@ namespace MiniT2I {
             int64_t H        = x->ne[1];
             int64_t img_side = H / config.patch_size;
             int64_t txt_len  = context->ne[1];
-            ensure_position_cache(img_side, txt_len);
+            sd::PECacheKey pos_key("minit2i.pos_embed");
+            pos_key.add(img_side).add(config.hidden_size);
+            auto pos_embed = make_input(pe_cache_.get_or_create(pos_key, [&]() {
+                return sd::Tensor<float>({config.hidden_size, img_side * img_side, 1},
+                                         make_2d_sincos_pos_embed(static_cast<int>(img_side), static_cast<int>(config.hidden_size)));
+            }));
+            sd::PECacheKey txt_key("minit2i.txt_pe");
+            txt_key.add(txt_len).add(config.head_dim);
+            auto txt_pe = make_input(pe_cache_.get_or_create(txt_key, [&]() {
+                return sd::Tensor<float>({2, 2, config.head_dim / 2, txt_len},
+                                         make_text_rope(static_cast<int>(txt_len), static_cast<int>(config.head_dim)));
+            }));
+            sd::PECacheKey joint_key("minit2i.joint_pe");
+            joint_key.add(img_side).add(txt_len).add(config.head_dim).add(circular_x_enabled).add(circular_y_enabled);
+            auto joint_pe = make_input(pe_cache_.get_or_create(joint_key, [&]() {
+                auto values = make_text_rope(static_cast<int>(txt_len), static_cast<int>(config.head_dim));
+                auto image  = finish_rope_pe(make_vision_rope(static_cast<int>(img_side), static_cast<int>(config.head_dim)));
+                values.insert(values.end(), image.begin(), image.end());
+                return sd::Tensor<float>({2, 2, config.head_dim / 2, txt_len + img_side * img_side}, std::move(values));
+            }));
 
             auto runner_ctx = get_context();
-            auto out        = model.forward(&runner_ctx, x, context, mask, cached_pos_embed, cached_txt_pe, cached_joint_pe);
+            auto out        = model.forward(&runner_ctx, x, context, mask, pos_embed, txt_pe, joint_pe);
             ggml_build_forward_expand(gf, out);
             return gf;
         }

src/model/diffusion/moss_image.hpp

@@ -3,6 +3,7 @@
 
 #include <algorithm>
 #include <cinttypes>
+#include "core/pe_cache.h"
 
 #include "model/common/ggml_block.hpp"
 #include "model/common/rope.hpp"
@@ -193,9 +194,10 @@ namespace MOSSImage {
     };
 
     struct MOSSImageRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
         MOSSImageConfig config;
         Model model;
-        std::vector<float> pe_data, time_data;
+        std::vector<float> time_data;
 
         MOSSImageRunner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
             : DiffusionModelRunner(backend, prefix, weight_manager), config(MOSSImageConfig::detect_from_weights(weights, prefix)), model(config) {
@@ -224,6 +226,9 @@ namespace MOSSImage {
                 auto text = make_input(*inputs.context);
                 GGML_ASSERT(x->ne[3] == 1 && x->ne[2] == config.channels && text->ne[0] == config.text_dim);
                 int64_t h = x->ne[1], w = x->ne[0], nt = text->ne[1];
+                sd::PECacheKey pe_key("moss_image.pe");
+                pe_key.add(h).add(w).add(nt).add(config.head_dim).add(circular_x_enabled).add(circular_y_enabled);
+                auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
                     Rope::Embedding embedding;
                     embedding.positions.append_image(static_cast<int>(h), static_cast<int>(w));
                     embedding.positions.images.back().height_axis = 2;
@@ -239,16 +244,16 @@ namespace MOSSImage {
                         embedding.ids.push_back({pos[0], pos[1], pos[2], pos[3]});
                     }
                     embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, std::vector<int>(4, static_cast<int>(config.head_dim / 4)), embedding.layout, &embedding.frequencies);
-                pe_data          = finish_rope_pe(std::move(embedding));
+                    auto values      = finish_rope_pe(std::move(embedding));
+                    return sd::Tensor<float>({2, 2, config.head_dim / 2, h * w + nt}, std::move(values));
+                }));
                 time_data.resize(config.freq_dim);
                 for (int64_t i = 0; i < config.freq_dim / 2; ++i) {
                     double angle                       = inputs.timesteps->values()[0] * std::pow(10000.0, -double(i) / double(config.freq_dim / 2));
                     time_data[i]                       = float(std::cos(angle));
                     time_data[i + config.freq_dim / 2] = float(std::sin(angle));
                 }
-                auto pe   = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, h * w + nt);
                 auto time = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.freq_dim, 1, 1);
-                set_backend_tensor_data(pe, pe_data.data());
                 set_backend_tensor_data(time, time_data.data());
                 auto ctx = get_context();
                 ggml_build_forward_expand(gf, model.forward(&ctx, x, time, text, pe, skip));

src/model/diffusion/pid.hpp

@@ -7,6 +7,7 @@
 #include <memory>
 #include <string>
 #include <vector>
+#include "core/pe_cache.h"
 
 #include "core/ggml_extend.h"
 #include "core/ggml_runner.h"
@@ -822,12 +823,9 @@ namespace Pid {
     };
 
     struct PiDRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
         PixelDiTConfig config;
         PixelDiT model;
-        std::vector<float> pos_img_vec;
-        std::vector<float> pos_txt_vec;
-        std::vector<float> pixel_pos_vec;
-        std::vector<float> pixel_pos_comp_vec;
 
         PiDRunner(ggml_backend_t backend,
                   const String2TensorStorage& tensor_storage_map,
@@ -867,57 +865,53 @@ namespace Pid {
             int64_t Hs = Hp / config.patch_size;
             int64_t Ws = Wp / config.patch_size;
 
-            pos_img_vec  = finish_rope_pe(make_rope_2d(static_cast<int>(Hs),
+            sd::PECacheKey pos_img_key("pid.pos_img");
+            pos_img_key.add(Hs).add(Ws).add(config.hidden_size).add(config.num_groups).add(config.rope_ref_grid_h);
+            pos_img_key.add(config.rope_ref_grid_w).add(circular_x_enabled).add(circular_y_enabled);
+            auto pos_img = make_input(pe_cache_.get_or_create(pos_img_key, [&]() {
+                auto 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,
-                                              2,
-                                              config.hidden_size / config.num_groups / 2,
-                                              Hs * Ws);
-            set_backend_tensor_data(pos_img, pos_img_vec.data());
+                return sd::Tensor<float>({2, 2, config.hidden_size / config.num_groups / 2, Hs * Ws}, std::move(pos_img_vec));
+            }));
 
             int64_t Ltxt = std::min<int64_t>(context->ne[1], config.txt_max_length);
-            pos_txt_vec  = make_rope_1d(static_cast<int>(Ltxt),
+            sd::PECacheKey pos_txt_key("pid.pos_txt");
+            pos_txt_key.add(Ltxt).add(config.hidden_size).add(config.num_groups).add(config.text_rope_theta);
+            auto pos_txt = make_input(pe_cache_.get_or_create(pos_txt_key, [&]() {
+                auto pos_txt_vec = make_rope_1d(static_cast<int>(Ltxt),
                                                 static_cast<int>(config.hidden_size / config.num_groups),
                                                 config.text_rope_theta);
-            auto pos_txt = ggml_new_tensor_4d(compute_ctx,
-                                              GGML_TYPE_F32,
-                                              2,
-                                              2,
-                                              config.hidden_size / config.num_groups / 2,
-                                              Ltxt);
-            set_backend_tensor_data(pos_txt, pos_txt_vec.data());
+                return sd::Tensor<float>({2, 2, config.hidden_size / config.num_groups / 2, Ltxt}, std::move(pos_txt_vec));
+            }));
 
-            pixel_pos_vec  = make_pixel_abs_pos(static_cast<int>(Hp),
+            sd::PECacheKey pixel_pos_key("pid.pixel_pos");
+            pixel_pos_key.add(Hp).add(Wp).add(config.pixel_hidden_size);
+            auto pixel_pos = make_input(pe_cache_.get_or_create(pixel_pos_key, [&]() {
+                auto pixel_pos_vec = make_pixel_abs_pos(static_cast<int>(Hp),
                                                         static_cast<int>(Wp),
                                                         static_cast<int>(config.pixel_hidden_size));
-            auto pixel_pos = ggml_new_tensor_3d(compute_ctx,
-                                                GGML_TYPE_F32,
-                                                config.pixel_hidden_size,
-                                                Wp * Hp,
-                                                1);
-            set_backend_tensor_data(pixel_pos, pixel_pos_vec.data());
+                return sd::Tensor<float>({config.pixel_hidden_size, Wp * Hp, 1}, std::move(pixel_pos_vec));
+            }));
 
-            pixel_pos_comp_vec  = finish_rope_pe(make_rope_2d(static_cast<int>(Hs),
+            sd::PECacheKey pixel_pos_comp_key("pid.pixel_pos_comp");
+            pixel_pos_comp_key.add(Hs).add(Ws).add(config.pixel_attn_hidden_size).add(config.pixel_num_groups);
+            pixel_pos_comp_key.add(config.rope_ref_grid_h).add(config.rope_ref_grid_w).add(circular_x_enabled);
+            pixel_pos_comp_key.add(circular_y_enabled);
+            auto pixel_pos_comp = make_input(pe_cache_.get_or_create(pixel_pos_comp_key, [&]() {
+                auto 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,
-                                                     2,
-                                                     config.pixel_attn_hidden_size / config.pixel_num_groups / 2,
-                                                     Hs * Ws);
-            set_backend_tensor_data(pixel_pos_comp, pixel_pos_comp_vec.data());
+                return sd::Tensor<float>({2, 2, config.pixel_attn_hidden_size / config.pixel_num_groups / 2, Hs * Ws}, std::move(pixel_pos_comp_vec));
+            }));
 
             auto runner_ctx = get_context();
             auto out        = model.forward(&runner_ctx,

src/model/diffusion/pixart.hpp

@@ -5,6 +5,7 @@
 #include <cinttypes>
 #include <cmath>
 #include <vector>
+#include "core/pe_cache.h"
 
 #include "core/ggml_extend.h"
 #include "core/ggml_runner.h"
@@ -295,9 +296,9 @@ namespace PixArt {
     };
 
     struct PixArtRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
         PixArtConfig config;
         PixArtModel model;
-        std::vector<float> pos_vec;
 
         PixArtRunner(ggml_backend_t backend,
                      const String2TensorStorage& tensor_storage_map      = {},
@@ -357,10 +358,13 @@ namespace PixArt {
             int64_t wp = W / config.patch_size;
             int64_t hp = H / config.patch_size;
 
-            pos_vec  = gen_2d_sincos_pos_embed(config.hidden_size, hp, wp,
+            sd::PECacheKey pos_key("pixart.pos");
+            pos_key.add(hp).add(wp).add(config.hidden_size).add(config.pos_embed_base_size).add(config.interpolation_scale);
+            auto pos = make_input(pe_cache_.get_or_create(pos_key, [&]() {
+                auto pos_vec = gen_2d_sincos_pos_embed(config.hidden_size, hp, wp,
                                                        config.pos_embed_base_size, config.interpolation_scale);
-            auto pos = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.hidden_size, wp * hp, 1);
-            set_backend_tensor_data(pos, pos_vec.data());
+                return sd::Tensor<float>({config.hidden_size, wp * hp, 1}, std::move(pos_vec));
+            }));
 
             auto runner_ctx  = get_context();
             ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pos, context_mask);

src/model/diffusion/qwen_image.hpp

@@ -4,6 +4,7 @@
 #include <memory>
 #include "core/ggml_extend_backend.h"
 #include "core/ggml_tensor_utils.h"
+#include "core/pe_cache.h"
 
 #include "core/util.h"
 #include "model/common/block.hpp"
@@ -561,10 +562,11 @@ namespace Qwen {
     };
 
     struct QwenImageRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
+
     public:
         QwenImageConfig config;
         QwenImageModel qwen_image;
-        std::vector<float> pe_vec;
         std::vector<float> modulate_index_vec;
         std::vector<int32_t> additional_t_cond_vec;
         SDVersion version;
@@ -635,7 +637,12 @@ namespace Qwen {
                 ref_index_mode = Rope::RefIndexMode::DECREASE;
             }
 
-            pe_vec      = finish_rope_pe(Rope::gen_qwen_image_pe(time_len,
+            sd::PECacheKey pe_key("qwen_image.pe");
+            pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+            pe_key.add(circular_x_enabled).add(circular_y_enabled).add(time_len).add(batch_size).add(config.patch_size);
+            pe_key.add_shapes(ref_latents).add(ref_index_mode);
+            auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+                auto 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,
@@ -646,12 +653,8 @@ namespace Qwen {
                                                                      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);
-            // pe->data = pe_vec.data();
-            // print_ggml_tensor(pe, true, "pe");
-            // pe->data = nullptr;
-            set_backend_tensor_data(pe, pe_vec.data());
+                return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+            }));
 
             ggml_tensor* modulate_index = nullptr;
             if (config.zero_cond_t) {

src/model/diffusion/qwen_image_2_1.hpp

@@ -1,6 +1,7 @@
 #ifndef __SD_MODEL_DIFFUSION_QWEN_IMAGE_2_1_H__
 #define __SD_MODEL_DIFFUSION_QWEN_IMAGE_2_1_H__
 
+#include "core/pe_cache.h"
 #include "model/diffusion/qwen_image.hpp"
 
 namespace Qwen {
@@ -370,9 +371,9 @@ namespace Qwen {
     };
 
     struct QwenImage21Runner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
         QwenImage21Config config;
         QwenImage21Model model;
-        std::vector<float> pe_data;
         std::vector<sd::Tensor<float>> mask_data;
         ggml_type prefix_cache_type = GGML_TYPE_COUNT;
         bool prefix_cache_enabled   = true;
@@ -505,19 +506,6 @@ namespace Qwen {
                     checked_cache.flash_attn_used = &flash_attn_used;
                 }
                 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);
-                    }
-                }
-                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) {
@@ -535,9 +523,30 @@ namespace Qwen {
                 }
                 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());
+                    sd::PECacheKey pe_key("qwen_image_2_1.pe");
+                    pe_key.add(layout.positions).add(cached).add(layout.prefix_length).add(config.axes_dim).add(config.head_dim);
+                    pe_key.add(circular_x_enabled).add(circular_y_enabled);
+                    for (const auto& region : layout.rope_layout.images) {
+                        pe_key.add(region.begin).add(region.count).add(region.height_period).add(region.width_period);
+                        pe_key.add(region.height_axis).add(region.width_axis);
+                    }
+                    auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+                        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);
+                            }
+                        }
+                        embedding.values     = Rope::embed_nd(embedding.ids, 1, 10000.f, config.axes_dim, embedding.layout, &embedding.frequencies);
+                        auto values          = finish_rope_pe(std::move(embedding));
+                        const int64_t tokens = static_cast<int64_t>(layout.positions.size()) - (cached ? layout.prefix_length : 0);
+                        return sd::Tensor<float>({2, 2, config.head_dim / 2, tokens}, std::move(values));
+                    }));
                     std::vector<ggml_tensor*> masks, ref_inputs;
                     for (const auto& mask : mask_data) {
                         masks.push_back(mask.empty() ? nullptr : make_input(mask));

src/model/diffusion/wan.hpp

@@ -8,6 +8,7 @@
 #include <utility>
 #include "core/ggml_extend_backend.h"
 #include "core/ggml_tensor_utils.h"
+#include "core/pe_cache.h"
 
 #include "model/common/block.hpp"
 #include "model/common/rope.hpp"
@@ -885,11 +886,12 @@ namespace WAN {
     };
 
     struct WanRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
+
     public:
         std::string desc = "wan";
         WanConfig config;
         Wan wan;
-        std::vector<float> pe_vec;
         SDVersion version;
 
         WanRunner(ggml_backend_t backend,
@@ -1000,7 +1002,11 @@ namespace WAN {
             ggml_tensor* audio_embed     = make_optional_input(audio_embed_tensor);
             ggml_tensor* ref_latent      = make_optional_input(ref_latent_tensor);
 
-            pe_vec = Rope::gen_wan_pe(static_cast<int>(x->ne[2]),
+            sd::PECacheKey pe_key("wan.pe");
+            pe_key.add_shape(x).add_shape(ref_latent).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+            pe_key.add(std::get<0>(config.patch_size)).add(std::get<1>(config.patch_size)).add(std::get<2>(config.patch_size));
+            auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+                auto pe_vec = Rope::gen_wan_pe(static_cast<int>(x->ne[2]),
                                                static_cast<int>(x->ne[1]),
                                                static_cast<int>(x->ne[0]),
                                                std::get<0>(config.patch_size),
@@ -1010,7 +1016,6 @@ namespace WAN {
                                                config.theta,
                                                config.axes_dim);
                 if (ref_latent != nullptr) {
-                // Match S2V's reference-frame temporal offset.
                     int t_start = std::max(30, static_cast<int>(x->ne[2]) + 9);
                     auto ref_pe = Rope::gen_wan_pe(static_cast<int>(ref_latent->ne[2]),
                                                    static_cast<int>(ref_latent->ne[1]),
@@ -1025,12 +1030,8 @@ namespace WAN {
                     pe_vec.insert(pe_vec.end(), ref_pe.begin(), ref_pe.end());
                 }
                 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);
-            // pe->data = pe_vec.data();
-            // print_ggml_tensor(pe);
-            // pe->data = nullptr;
-            set_backend_tensor_data(pe, pe_vec.data());
+                return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+            }));
 
             if (c_concat != nullptr) {
                 x = ggml_concat(compute_ctx, x, c_concat, 3);

src/model/diffusion/z_image.hpp

@@ -3,6 +3,7 @@
 
 #include <algorithm>
 #include <cinttypes>
+#include "core/pe_cache.h"
 
 #include "core/ggml_extend.h"
 #include "core/ggml_extend_backend.h"
@@ -916,10 +917,11 @@ namespace ZImage {
     };
 
     struct ZImageRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
+
     public:
         ZImageConfig config;
         ZImageModel z_image;
-        std::vector<float> pe_vec;
         std::vector<float> timestep_vec;
         SDVersion version;
 
@@ -971,7 +973,12 @@ namespace ZImage {
 
             ggml_tensor* control_latent = make_optional_input(control_latent_tensor);
 
-            pe_vec      = finish_rope_pe(Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
+            sd::PECacheKey pe_key("z_image.pe");
+            pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+            pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size).add_shapes(ref_latents);
+            pe_key.add(ref_index_mode);
+            auto pe         = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+                auto 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]),
@@ -982,12 +989,8 @@ namespace ZImage {
                                                                   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);
-            // pe->data = pe_vec.data();
-            // print_ggml_tensor(pe, true, "pe");
-            // pe->data = nullptr;
-            set_backend_tensor_data(pe, pe_vec.data());
+                return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+            }));
             auto runner_ctx = get_context();
 
             ggml_tensor* out = z_image.forward(&runner_ctx,

src/model/diffusion/z_image_l2p.hpp

@@ -2,6 +2,7 @@
 #define __SD_MODEL_DIFFUSION_Z_IMAGE_L2P_HPP__
 
 #include <cmath>
+#include "core/pe_cache.h"
 
 #include "z_image.hpp"
 
@@ -183,9 +184,9 @@ namespace ZImageL2P {
     };
 
     struct ZImageL2PRunner : public DiffusionModelRunner {
+        sd::PECache pe_cache_;
         ZImageL2PConfig config;
         ZImageL2PModel model;
-        std::vector<float> pe_vec;
 
         ZImageL2PRunner(ggml_backend_t backend,
                         const String2TensorStorage& tensors,
@@ -221,7 +222,11 @@ namespace ZImageL2P {
                 auto context = make_input(*inputs.context);
                 GGML_ASSERT(x->ne[3] == 1);
                 GGML_ASSERT(x->ne[0] % config.patch_size == 0 && x->ne[1] % config.patch_size == 0);
-                pe_vec      = finish_rope_pe(Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
+                sd::PECacheKey pe_key("z_image_l2p.pe");
+                pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum);
+                pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size);
+                auto pe  = make_input(pe_cache_.get_or_create(pe_key, [&]() {
+                    auto 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]),
@@ -232,8 +237,8 @@ namespace ZImageL2P {
                                                                       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());
+                    return sd::Tensor<float>({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec));
+                }));
                 auto ctx = get_context();
                 auto out = model.forward(&ctx, x, t, context, pe);
                 ggml_build_forward_expand(gf, out);