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

+19 -14
View File
@@ -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,19 +648,23 @@ 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),
static_cast<int>(H),
static_cast<int>(W),
config.patch_t,
config.patch_h,
config.patch_w,
static_cast<int>(N),
static_cast<int>(context->ne[1]),
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());
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,
config.patch_h,
config.patch_w,
static_cast<int>(N),
static_cast<int>(context->ne[1]),
config.rope_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 = lingbot_video.forward(&runner_ctx, x, timesteps, context, pe);