upstream-archive/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-09-29 01:18:05 -05:00
Code Issues Packages Projects Releases Wiki Activity

feat: add configurable Qwen cache types and early cache scheduling (#2045)

Browse Source
This commit is contained in:
leejet
2026-09-25 01:51:01 +08:00
committed by GitHub
parent 4dfe8f5d45
commit 740c7ae193
11 changed files with 185 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/core/ggml_extend.cpp
+8 -1
View File
@@ -623,7 +623,11 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
bool skip_reshape,
bool flash_attn,
float kv_scale,
bool sage_attn) { // avoid overflow
bool sage_attn,
bool* used_flash_attn) { // avoid overflow
if (used_flash_attn != nullptr) {
*used_flash_attn = false;
}
int64_t L_q;
int64_t L_k;
int64_t C;
@@ -755,6 +759,9 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
if (can_use_flash_attn) {
kqv = build_kqv(q, k, v, mask);
if (kqv != nullptr) {
if (used_flash_attn != nullptr) {
*used_flash_attn = true;
}
kqv = ggml_view_4d(ctx,
kqv,
d_head,
src/core/ggml_extend.h
+6 -5
View File
@@ -217,11 +217,12 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
ggml_tensor* k,
ggml_tensor* v,
int64_t n_head,
ggml_tensor* mask = nullptr,
bool skip_reshape = false,
bool flash_attn = false,
float kv_scale = 1.0f,
bool sage_attn = false);
ggml_tensor* mask = nullptr,
bool skip_reshape = false,
bool flash_attn = false,
float kv_scale = 1.0f,
bool sage_attn = false,
bool* used_flash_attn = nullptr);
ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx,
ggml_tensor* x,
src/core/ggml_runner.cpp
+12 -6
View File
@@ -21,11 +21,12 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
ggml_tensor* mask,
bool skip_reshape,
bool flash_attn,
float kv_scale) {
float kv_scale,
bool* used_flash_attn) {
if (ctx->attn_scale > 0.f) {
kv_scale = ctx->attn_scale;
}
return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale, ctx->sage_attn_enabled);
return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale, ctx->sage_attn_enabled, used_flash_attn);
}
void GGMLRunner::alloc_params_ctx() {
@@ -515,9 +516,10 @@ GGMLRunner::~GGMLRunner() {
free_params_ctx();
}
GGMLRunnerContext GGMLRunner::get_context() {
GGMLRunnerContext GGMLRunner::get_context(ggml_cgraph* graph) {
GGMLRunnerContext runner_ctx;
runner_ctx.ggml_ctx = compute_ctx;
runner_ctx.graph = graph;
runner_ctx.backend = runtime_backend;
runner_ctx.flash_attn_enabled = flash_attn_enabled;
runner_ctx.sage_attn_enabled = sage_attn_enabled;
@@ -532,8 +534,8 @@ GGMLRunnerContext GGMLRunner::get_context() {
runner_ctx.get_cache_tensor = [this](const std::string& name) {
return this->get_cache_tensor_by_name(name);
};
runner_ctx.cache_tensor = [this](const std::string& name, ggml_tensor* tensor) {
this->cache(name, tensor);
runner_ctx.cache_tensor = [this, graph](const std::string& name, ggml_tensor* tensor) {
this->cache(name, tensor, graph);
};
runner_ctx.set_backend_tensor_data = [this](ggml_tensor* tensor, const void* data) {
this->set_backend_tensor_data(tensor, data);
@@ -575,7 +577,7 @@ ggml_tensor* GGMLRunner::to_backend(ggml_tensor* tensor) {
}
}
void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) {
void GGMLRunner::cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* graph) {
if (tensor != nullptr && tensor->view_src != nullptr) {
tensor = ggml_cont(compute_ctx, tensor);
}
@@ -583,6 +585,10 @@ void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) {
ggml_set_output(tensor);
}
cache_.stage(name, tensor);
if (graph != nullptr && tensor != nullptr) {
// Schedule the cache output here so its source can be reused before graph end.
ggml_build_forward_expand(graph, tensor);
}
}
std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
src/core/ggml_runner.h
+15 -6
View File
@@ -67,6 +67,7 @@ struct WeightAdapter {
struct GGMLRunnerContext {
ggml_backend_t backend = nullptr;
ggml_context* ggml_ctx = nullptr;
ggml_cgraph* graph = nullptr;
bool flash_attn_enabled = false;
bool sage_attn_enabled = false;
float linear_scale = 0.f;
@@ -102,6 +103,12 @@ struct GGMLRunnerContext {
return get_cache_tensor(name);
}
void expand_graph(ggml_tensor* tensor) const {
if (graph != nullptr && tensor != nullptr) {
ggml_build_forward_expand(graph, tensor);
}
}
void persist_cache_tensor(const std::string& name, ggml_tensor* tensor) const {
if (!cache_tensor || tensor == nullptr) {
return;
@@ -122,10 +129,11 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
ggml_tensor* k,
ggml_tensor* v,
int64_t n_head,
ggml_tensor* mask = nullptr,
bool skip_reshape = false,
bool flash_attn = false,
float kv_scale = 1.f);
ggml_tensor* mask = nullptr,
bool skip_reshape = false,
bool flash_attn = false,
float kv_scale = 1.f,
bool* used_flash_attn = nullptr);
struct GGMLRunner {
private:
@@ -289,7 +297,8 @@ public:
virtual ~GGMLRunner();
virtual GGMLRunnerContext get_context();
// Binding a graph schedules cache outputs at registration instead of graph end.
virtual GGMLRunnerContext get_context(ggml_cgraph* graph = nullptr);
void reset_compute_ctx();
@@ -324,7 +333,7 @@ public:
ggml_tensor* to_backend(ggml_tensor* tensor);
void cache(const std::string name, ggml_tensor* tensor);
void cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* graph = nullptr);
ggml_tensor* get_cache_tensor_by_name(const std::string& name) {
return cache_.get(name);
src/model/diffusion/control.hpp
+2 -3
View File
@@ -376,7 +376,7 @@ struct ControlNet : public GGMLRunner {
hint = make_input(hint_tensor);
}
auto runner_ctx = get_context();
auto runner_ctx = get_context(gf);
auto outs = control_net.forward(&runner_ctx,
x,
@@ -389,8 +389,7 @@ struct ControlNet : public GGMLRunner {
if (guided_hint_input == nullptr && !outs.empty()) {
guided_hint_output_ggml = outs[0];
ggml_set_output(guided_hint_output_ggml);
cache(guided_hint_cache_name(), guided_hint_output_ggml);
ggml_build_forward_expand(gf, guided_hint_output_ggml);
runner_ctx.persist_cache_tensor(guided_hint_cache_name(), guided_hint_output_ggml);
}
control_outputs_ggml.reserve(outs.size() > 0 ? outs.size() - 1 : 0);
src/model/diffusion/qwen_image_2_1.hpp
+109 -19
View File
@@ -134,6 +134,8 @@ namespace Qwen {
std::string name;
std::string cut_group;
int64_t prefix_length = 0;
ggml_type type = GGML_TYPE_F32;
bool* flash_attn_used = nullptr;
};
class QwenImage21ZeroCenterRMSNorm : public RMSNorm {
@@ -189,10 +191,15 @@ namespace Qwen {
q = Rope::apply_rope(ctx->ggml_ctx, q, pe);
k = Rope::apply_rope(ctx->ggml_ctx, k, pe);
if (cache.mode == QwenImage21PrefixCache::Mode::STORE) {
// Preserve query-first attention evaluation while writing each layer's
// prefix before its full-sequence K/V can accumulate across layers.
ctx->expand_graph(q);
auto persist = [&](ggml_tensor* tensor, int axis, const char* name) {
auto part = ggml_ext_slice(ctx->ggml_ctx, tensor, axis, 0, cache.prefix_length);
auto copy = ggml_new_tensor(ctx->ggml_ctx, GGML_TYPE_F32, 4, part->ne);
copy = ggml_cpy(ctx->ggml_ctx, part, copy);
// Pack the contiguous data into wider rows so quantization blocks
// can exceed head_dim without padding or changing element order.
part = ggml_reshape_2d(ctx->ggml_ctx, part, x->ne[0], cache.prefix_length);
auto copy = ggml_cast(ctx->ggml_ctx, part, cache.type);
// Keep the copy in this layer's segment so graph cuts do not
// retain or recompute the full-sequence K/V in the final segment.
sd::ggml_graph_cut::mark_graph_cut(copy, cache.cut_group, name);
@@ -201,21 +208,37 @@ namespace Qwen {
persist(k, 1, "k");
persist(v, 2, "v");
}
auto attend = [&](ggml_tensor* aq, ggml_tensor* ak, ggml_tensor* av, ggml_tensor* mask) {
bool used_flash_attn = false;
auto out = ggml_ext_attention_ext(ctx, aq, ak, av, heads, mask, true, ctx->flash_attn_enabled, 1.f, &used_flash_attn);
if (cache.flash_attn_used != nullptr) {
*cache.flash_attn_used &= used_flash_attn;
}
return out;
};
ggml_tensor* result = nullptr;
if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) {
auto prefix_k = ctx->load_cache_tensor(cache.name + ".k");
auto prefix_v = ctx->load_cache_tensor(cache.name + ".v");
GGML_ASSERT(prefix_k != nullptr && prefix_v != nullptr);
k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 1);
v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2);
result = ggml_ext_attention_ext(ctx, q, k, v, heads, nullptr, true, ctx->flash_attn_enabled);
if (prefix_k->type != k->type) {
prefix_k = ggml_cast(ctx->ggml_ctx, prefix_k, k->type);
}
if (prefix_v->type != v->type) {
prefix_v = ggml_cast(ctx->ggml_ctx, prefix_v, v->type);
}
prefix_k = ggml_reshape_4d(ctx->ggml_ctx, prefix_k, dim_head, cache.prefix_length, heads, k->ne[3]);
prefix_v = ggml_reshape_4d(ctx->ggml_ctx, prefix_v, dim_head, heads, cache.prefix_length, v->ne[3]);
k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 1);
v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2);
result = attend(q, k, v, nullptr);
} else {
for (size_t i = 0; i < segments.size(); ++i) {
const auto& segment = segments[i];
auto sq = ggml_ext_slice(ctx->ggml_ctx, q, 1, segment.start, segment.end);
auto sk = ggml_ext_slice(ctx->ggml_ctx, k, 1, 0, segment.end);
auto sv = ggml_ext_slice(ctx->ggml_ctx, v, 2, 0, segment.end);
auto out = ggml_ext_attention_ext(ctx, sq, sk, sv, heads, masks[i], true, ctx->flash_attn_enabled);
auto out = attend(sq, sk, sv, masks[i]);
result = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1);
}
}
@@ -351,8 +374,35 @@ namespace Qwen {
QwenImage21Model model;
std::vector<float> pe_data;
std::vector<sd::Tensor<float>> mask_data;
bool prefix_cache_enabled = true;
bool prefix_cache_disabled = false;
ggml_type prefix_cache_type = GGML_TYPE_COUNT;
bool prefix_cache_enabled = true;
bool prefix_cache_disabled = false;
bool prefix_cache_auto_f32 = false;
static bool supports_prefix_cache_type(ggml_type type) {
if (type == GGML_TYPE_F32) {
return true;
}
const auto* traits = ggml_get_type_traits(type);
if (traits->from_float_ref == nullptr || traits->to_float == nullptr) {
return false;
}
auto cpu = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
if (cpu == nullptr) {
return false;
}
auto ctx = std::unique_ptr<ggml_context, decltype(&ggml_free)>(
ggml_init({3 * ggml_tensor_overhead(), nullptr, true}), ggml_free);
if (ctx == nullptr) {
return false;
}
// Some reference quantizers have no runtime copy support. Query the
// device through the registry so dynamically loaded CPU backends work.
auto source = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, ggml_blck_size(type));
auto encoded = ggml_cast(ctx.get(), source, type);
auto decoded = ggml_cast(ctx.get(), encoded, GGML_TYPE_F32);
return ggml_backend_dev_supports_op(cpu, encoded) && ggml_backend_dev_supports_op(cpu, decoded);
}
QwenImage21Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr, const char* model_args = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager),
@@ -361,6 +411,22 @@ namespace Qwen {
for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
if (key == "qwen_image_2_1_prefix_cache" && !parse_strict_bool(value, prefix_cache_enabled)) {
LOG_WARN("ignoring invalid Qwen Image 2.1 model arg '%s=%s'", key.c_str(), value.c_str());
} else if (key == "qwen_image_2_1_prefix_cache_type") {
if (value == "auto") {
prefix_cache_type = GGML_TYPE_COUNT;
continue;
}
const auto type = sd_type_to_ggml_type(str_to_sd_type(value.c_str()));
if (type == GGML_TYPE_COUNT) {
LOG_WARN("ignoring unknown Qwen Image 2.1 cache type '%s'", value.c_str());
} else if (!supports_prefix_cache_type(type)) {
LOG_WARN("ignoring Qwen Image 2.1 cache type '%s': runtime conversion to and from F32 is unavailable", value.c_str());
} else if (config.hidden_size % ggml_blck_size(type) != 0) {
LOG_WARN("ignoring Qwen Image 2.1 cache type '%s': block size %" PRId64 " does not divide hidden size %" PRId64,
value.c_str(), ggml_blck_size(type), config.hidden_size);
} else {
prefix_cache_type = type;
}
}
}
model.init(params_ctx, weights, prefix);
@@ -377,11 +443,9 @@ namespace Qwen {
const auto name = cache.name + "." + std::to_string(i);
auto k = get_cache_tensor_by_name(name + ".k");
auto v = get_cache_tensor_by_name(name + ".v");
if (k == nullptr || v == nullptr || k->type != GGML_TYPE_F32 || v->type != GGML_TYPE_F32 ||
k->ne[0] != config.head_dim || k->ne[1] != cache.prefix_length ||
k->ne[2] != config.hidden_size / config.head_dim || k->ne[3] != 1 ||
v->ne[0] != config.head_dim || v->ne[1] != config.hidden_size / config.head_dim ||
v->ne[2] != cache.prefix_length || v->ne[3] != 1) {
if (k == nullptr || v == nullptr || k->type != cache.type || v->type != cache.type ||
k->ne[0] != config.hidden_size || k->ne[1] != cache.prefix_length || k->ne[2] != 1 || k->ne[3] != 1 ||
v->ne[0] != config.hidden_size || v->ne[1] != cache.prefix_length || v->ne[2] != 1 || v->ne[3] != 1) {
return false;
}
}
@@ -418,15 +482,28 @@ namespace Qwen {
}
if (!runner_started()) {
prefix_cache_disabled = false;
prefix_cache_auto_f32 = false;
}
QwenImage21PrefixCache cache;
if (prefix_cache_enabled && !prefix_cache_disabled && extra != nullptr && extra->prefix_id != 0 && layout.prefix_length > 0) {
cache.name = "qwen_image_2_1.prefix." + std::to_string(extra->prefix_id) +
".circular." + std::to_string(circular_x_enabled) + std::to_string(circular_y_enabled);
cache.prefix_length = layout.prefix_length;
cache.mode = has_prefix_cache(cache) ? QwenImage21PrefixCache::Mode::REUSE : QwenImage21PrefixCache::Mode::STORE;
if (prefix_cache_type != GGML_TYPE_COUNT) {
cache.type = prefix_cache_type;
} else if (!prefix_cache_auto_f32 && flash_attn_enabled && !sage_attn_enabled &&
(attn_scale <= 0.f || attn_scale == 1.f)) {
cache.type = GGML_TYPE_F16;
}
cache.mode = has_prefix_cache(cache) ? QwenImage21PrefixCache::Mode::REUSE : QwenImage21PrefixCache::Mode::STORE;
}
auto run = [&](const QwenImage21PrefixCache& active_cache) {
bool flash_attn_used = true;
auto run = [&](const QwenImage21PrefixCache& active_cache) {
flash_attn_used = true;
auto checked_cache = active_cache;
if (prefix_cache_type == GGML_TYPE_COUNT && active_cache.type == GGML_TYPE_F16) {
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;
@@ -459,7 +536,7 @@ 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));
layout.positions.size() - (cached ? layout.prefix_length : 0));
set_backend_tensor_data(pe, pe_data.data());
std::vector<ggml_tensor*> masks, ref_inputs;
for (const auto& mask : mask_data) {
@@ -470,15 +547,28 @@ namespace Qwen {
ref_inputs.push_back(make_input(ref));
}
}
auto ctx = get_context();
auto ctx = get_context(graph);
auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), cached ? nullptr : make_input(context),
ref_inputs, pe, layout, masks, active_cache);
ref_inputs, pe, layout, masks, checked_cache);
if (!flash_attn_used) {
return static_cast<ggml_cgraph*>(nullptr);
}
ggml_build_forward_expand(graph, out);
return graph;
};
return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), x.dim());
};
auto result = run(cache);
if (result.empty() && !flash_attn_used) {
// Casting an F16 cache back to F32 cannot recover its original values.
// Recompute the prefix before executing a graph that falls back from FA.
free_cache_ctx_and_buffer();
prefix_cache_auto_f32 = true;
cache.type = GGML_TYPE_F32;
cache.mode = QwenImage21PrefixCache::Mode::STORE;
LOG_DEBUG("Qwen Image 2.1: Flash Attention unavailable; using F32 prefix caching for this sampling run");
result = run(cache);
}
if (result.empty() && last_compute_status() == GGML_STATUS_ALLOC_FAILED &&
(cache.mode != QwenImage21PrefixCache::Mode::NONE || !cache_.empty())) {
// The failed graph has ended before persistent inputs are released.
@@ -493,7 +583,7 @@ namespace Qwen {
prefix_cache_disabled = true;
LOG_WARN("Qwen Image 2.1: incomplete prefix cache; disabling it for this sampling run");
} else {
LOG_DEBUG("Qwen Image 2.1: cached prefix %" PRIu64 " (%" PRId64 " tokens)", extra->prefix_id, layout.prefix_length);
LOG_DEBUG("Qwen Image 2.1: cached prefix %" PRIu64 " (%" PRId64 " tokens, %s)", extra->prefix_id, layout.prefix_length, ggml_type_name(cache.type));
}
}
return result;
src/model/diffusion/sensenova_u1.h
+4 -11
View File
@@ -442,16 +442,9 @@ namespace SenseNovaU1 {
k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 2);
v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2);
} else {
// Keep dedicated graph outputs alive until the runner copies them
// into its persistent cache buffer after graph execution.
auto cache_k = ggml_dup_tensor(ctx->ggml_ctx, k);
cache_k = ggml_cpy(ctx->ggml_ctx, k, cache_k);
ggml_set_output(cache_k);
auto cache_v = ggml_dup_tensor(ctx->ggml_ctx, v);
cache_v = ggml_cpy(ctx->ggml_ctx, v, cache_v);
ggml_set_output(cache_v);
ctx->persist_cache_tensor(layer_cache + ".k", cache_k);
ctx->persist_cache_tensor(layer_cache + ".v", cache_v);
ctx->expand_graph(q);
ctx->persist_cache_tensor(layer_cache + ".k", k);
ctx->persist_cache_tensor(layer_cache + ".v", v);
}
q = ggml_cont(ctx->ggml_ctx,
@@ -687,7 +680,7 @@ namespace SenseNovaU1 {
ggml_set_name(attention_mask, "snu15.prefix.attention_mask");
set_backend_tensor_data(attention_mask, attention_mask_vec.data());
auto runner_ctx = get_context();
auto runner_ctx = get_context(graph);
auto text_model = model.text_model();
auto hidden = text_model->embed(&runner_ctx, ids);
hidden = text_model->forward(&runner_ctx,
src/model/vae/ltx_vae.hpp
+2 -3
View File
@@ -1300,7 +1300,7 @@ struct LTXVideoVAE : public VAE {
feat_map[feat_idx] = get_cache_tensor_by_name(temporal_feat_cache_name(feat_idx));
}
auto runner_ctx = get_context();
auto runner_ctx = get_context(gf);
int feat_count = 0;
ggml_tensor* out = vae.decode_tiled_chunk(&runner_ctx,
z,
@@ -1313,8 +1313,7 @@ struct LTXVideoVAE : public VAE {
for (int feat_idx = 0; feat_idx < feat_count && feat_idx < static_cast<int>(feat_map.size()); ++feat_idx) {
ggml_tensor* feat_cache = feat_map[static_cast<size_t>(feat_idx)];
if (feat_cache != nullptr) {
cache(temporal_feat_cache_name(static_cast<size_t>(feat_idx)), feat_cache);
ggml_build_forward_expand(gf, feat_cache);
runner_ctx.persist_cache_tensor(temporal_feat_cache_name(static_cast<size_t>(feat_idx)), feat_cache);
}
}
src/model/vae/wan_vae.hpp
+2 -3
View File
@@ -1443,15 +1443,14 @@ namespace WAN {
ggml_tensor* z = make_input(z_tensor);
auto runner_ctx = get_context();
auto runner_ctx = get_context(gf);
ggml_tensor* out = ae.decode_tiled_chunk(&runner_ctx, z, chunk_idx);
for (size_t feat_idx = 0; feat_idx < ae._feat_map.size(); feat_idx++) {
ggml_tensor* feat_cache = ae._feat_map[feat_idx];
if (feat_cache != nullptr) {
cache("feat_idx:" + std::to_string(feat_idx), feat_cache);
ggml_build_forward_expand(gf, feat_cache);
runner_ctx.persist_cache_tensor("feat_idx:" + std::to_string(feat_idx), feat_cache);
}
}