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

This commit is contained in:
leejet authored and GitHub committed 2026-10-25 05:18:30 +08:00
1 parent a85db96aa5
commit ea209dbfcf
31 files changed
+855 -610

No files matched your search

+16 -11
View File
@@ -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);