perf: cache CPU position embeddings across DiT runners

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,36 +660,40 @@ 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,
-                                                                          static_cast<int>(h_pad),
-                                                                          static_cast<int>(w_pad),
-                                                                          static_cast<int>(config.patch_size),
-                                                                          config.theta,
-                                                                          config.axes_dim,
-                                                                          4.0f,
-                                                                          4.0f,
-                                                                          1.0f,
-                                                                          ref_latents));
-            int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2));
-            auto image_pe         = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, image_pos_len);
-            set_backend_tensor_data(image_pe, image_pe_vec.data());
+            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),
+                                                                              config.theta,
+                                                                              config.axes_dim,
+                                                                              4.0f,
+                                                                              4.0f,
+                                                                              1.0f,
+                                                                              ref_latents));
+                int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2));
+                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;
             if (t5_ids != nullptr) {
                 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,17 +800,21 @@ 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]),
-                                                      static_cast<int>(x->ne[0]),
-                                                      config.patch_size,
-                                                      static_cast<int>(x->ne[3]),
-                                                      static_cast<int>(context->ne[1]),
-                                                      ref_latents,
-                                                      config.theta,
-                                                      config.axes_dim));
-            int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
-            auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
-            set_backend_tensor_data(pe, pe_vec.data());
+            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]),
+                                                          static_cast<int>(context->ne[1]),
+                                                          ref_latents,
+                                                          config.theta,
+                                                          config.axes_dim));
+                int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+                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,16 +416,20 @@ 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]),
-                                                                  static_cast<int>(x->ne[0]),
-                                                                  config.patch_size,
-                                                                  static_cast<int>(x->ne[3]),
-                                                                  static_cast<int>(context->ne[1]),
-                                                                  config.theta,
-                                                                  config.axes_dim));
-            int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
-            auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, config.axes_dim_sum, 1, pos_len, 2);
-            set_backend_tensor_data(pe, pe_vec.data());
+            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]),
+                                                                      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);
+                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);