DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-11 19:36:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/model/diffusion/pid.hpp
+45 -51
View File
@@ -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),
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());
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)));
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),
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());
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);
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),
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());
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));
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),
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());
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)));
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,