DarthAffe/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-09-24 20:20:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add Qwen Image 2.1 prefix KV cache (#2035)

Browse Source
This commit is contained in:
leejet 2026-09-23 23:08:21 +08:00 committed by GitHub
parent e6281b6318
commit 2dc7f5408a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 237 additions and 81 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
docs/qwen_image_2.1.md
View File

@ -40,6 +40,14 @@ Pass the reference image with `-r` and describe the edit in `-p`. Vision weights
For multiple reference images, repeat `-r` in the desired order, for example `-r first.png -r second.png`. For multiple reference images, repeat `-r` in the desired order, for example `-r first.png -r second.png`.
### Prefix cache
By default, the first denoising call for each fixed condition saves the text and reference-image keys and values from every transformer layer. Later calls only compute the target-image tokens. Positive and negative conditions use separate caches, which are released when sampling ends.
The cache uses FP32 on all attention backends. For the default 32-layer model, a prefix of 4096 tokens takes about 4 GiB per condition, in addition to weights and working buffers. The runner accounts for the cache when checking the memory budget. If a cached execution runs out of memory, it releases the prefix caches, disables caching for the rest of that sampling run, and retries the full sequence once. Per-step conditioning extensions currently use the full-sequence path.
Disable this optimization with `--model-args qwen_image_2_1_prefix_cache=false`. It reuses step-independent activations; numerical results can still differ slightly because the matrix sizes change.
### Alpha channel ### Alpha channel
This model supports alpha channel output. As the model determines whether to output a regular image or with transparency through the prompt, according to [official recommendation](https://github.com/QwenLM/Qwen-Image-2.1#transparent-image-generation-rgba), use the following prompt format for better results: This model supports alpha channel output. As the model determines whether to output a regular image or with transparency through the prompt, according to [official recommendation](https://github.com/QwenLM/Qwen-Image-2.1#transparent-image-generation-rgba), use the following prompt format for better results:

2
examples/common/common.cpp
View File

@ -518,7 +518,7 @@ ArgOptions SDContextParams::get_options() {
{"", {"",
"--model-args", "--model-args",
"extra model args, key=value list. Supports chroma_use_dit_mask, chroma_use_t5_mask, " "extra model args, key=value list. Supports chroma_use_dit_mask, chroma_use_t5_mask, "
"chroma_t5_mask_pad, qwen_image_zero_cond_t", "chroma_t5_mask_pad, qwen_image_zero_cond_t, qwen_image_2_1_prefix_cache",
(int)',', (int)',',
&model_args}, &model_args},
{"", {"",

18
src/core/ggml_runner.cpp
View File

@ -644,6 +644,10 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
std::optional<sd::Tensor<float>> output; std::optional<sd::Tensor<float>> output;
try { try {
output = execute_graph(graph, n_threads, no_return, read_outputs); output = execute_graph(graph, n_threads, no_return, read_outputs);
} catch (const std::bad_alloc&) {
last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
LOG_ERROR("%s graph allocation failed", get_desc().c_str());
return std::nullopt;
} catch (const std::exception& error) { } catch (const std::exception& error) {
last_compute_status_ = GGML_STATUS_FAILED; last_compute_status_ = GGML_STATUS_FAILED;
LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(), LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(),
@ -964,10 +968,16 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
} }
LOG_DEBUG("%s executing segment %zu/%zu: %s", get_desc().c_str(), LOG_DEBUG("%s executing segment %zu/%zu: %s", get_desc().c_str(),
index + 1, plan.segments.size(), segment.group_name.c_str()); index + 1, plan.segments.size(), segment.group_name.c_str());
if (!execute_segment(segment_graph, n_threads) || if (!execute_segment(segment_graph, n_threads)) {
!cache_.capture(segment_graph) || return fail_segment("execution");
!cut_cache_.capture(graph, segment, get_desc().c_str())) { }
return fail_segment("execution or output caching"); auto cache_status = cache_.capture(segment_graph);
if (cache_status == GGML_STATUS_SUCCESS) {
cache_status = cut_cache_.capture(graph, segment, get_desc().c_str());
}
if (cache_status != GGML_STATUS_SUCCESS) {
last_compute_status_ = cache_status;
return fail_segment("output caching");
} }
sync_runtime_residency(); sync_runtime_residency();
if (last) { if (last) {

30
src/core/runner_cache.cpp
View File

@ -26,10 +26,13 @@ namespace sd {
std::unique_ptr<CachedTensor> CachedTensor::copy(ggml_backend_t backend, std::unique_ptr<CachedTensor> CachedTensor::copy(ggml_backend_t backend,
const std::string& name, const std::string& name,
ggml_tensor* source) { ggml_tensor* source,
ggml_status& status) {
status = GGML_STATUS_FAILED;
if (ggml_graph_cut::tensor_buffer(source) == nullptr) { if (ggml_graph_cut::tensor_buffer(source) == nullptr) {
return nullptr; return nullptr;
} }
status = GGML_STATUS_ALLOC_FAILED;
auto entry = std::make_unique<CachedTensor>(); auto entry = std::make_unique<CachedTensor>();
entry->context = ggml_init({2 * ggml_tensor_overhead(), nullptr, true}); entry->context = ggml_init({2 * ggml_tensor_overhead(), nullptr, true});
if (entry->context == nullptr) { if (entry->context == nullptr) {
@ -50,6 +53,7 @@ namespace sd {
} else { } else {
ggml_backend_tensor_copy(source, entry->tensor); ggml_backend_tensor_copy(source, entry->tensor);
} }
status = GGML_STATUS_SUCCESS;
return entry; return entry;
} }
@ -106,9 +110,9 @@ namespace sd {
return pending > SIZE_MAX - committed ? SIZE_MAX : committed + pending; return pending > SIZE_MAX - committed ? SIZE_MAX : committed + pending;
} }
bool RunnerCache::capture(ggml_cgraph* graph) { ggml_status RunnerCache::capture(ggml_cgraph* graph) {
if (outputs_.empty()) { if (outputs_.empty()) {
return true; return GGML_STATUS_SUCCESS;
} }
const auto tensors = cache_graph_tensors(graph); const auto tensors = cache_graph_tensors(graph);
for (const auto& output : outputs_) { for (const auto& output : outputs_) {
@ -116,14 +120,15 @@ namespace sd {
continue; continue;
} }
GGML_ASSERT(ggml_is_contiguous(output.second)); GGML_ASSERT(ggml_is_contiguous(output.second));
auto entry = CachedTensor::copy(backend_, output.first, output.second); ggml_status status;
auto entry = CachedTensor::copy(backend_, output.first, output.second, status);
if (entry == nullptr) { if (entry == nullptr) {
return false; return status;
} }
pending_[output.first] = std::move(entry); pending_[output.first] = std::move(entry);
} }
ggml_backend_synchronize(backend_); ggml_backend_synchronize(backend_);
return true; return GGML_STATUS_SUCCESS;
} }
void RunnerCache::graph_end(bool success) { void RunnerCache::graph_end(bool success) {
@ -180,9 +185,9 @@ namespace sd {
} }
} }
bool GraphCutTensorCache::capture(ggml_cgraph* graph, ggml_status GraphCutTensorCache::capture(ggml_cgraph* graph,
const ggml_graph_cut::Segment& segment, const ggml_graph_cut::Segment& segment,
const char* log_desc) { const char* log_desc) {
size_t copied_bytes = 0; size_t copied_bytes = 0;
size_t copied_count = 0; size_t copied_count = 0;
for (int index : segment.output_node_indices) { for (int index : segment.output_node_indices) {
@ -191,10 +196,11 @@ namespace sd {
!segment.future_cut_names.count(output->name)) { !segment.future_cut_names.count(output->name)) {
continue; continue;
} }
auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output)); ggml_status status;
auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output), status);
if (entry == nullptr) { if (entry == nullptr) {
LOG_ERROR("%s failed to capture graph cut tensor: %s", log_desc, output->name); LOG_ERROR("%s failed to capture graph cut tensor: %s", log_desc, output->name);
return false; return status;
} }
const size_t size = ggml_backend_buffer_get_size(entry->buffer); const size_t size = ggml_backend_buffer_get_size(entry->buffer);
copied_bytes = size > SIZE_MAX - copied_bytes ? SIZE_MAX : copied_bytes + size; copied_bytes = size > SIZE_MAX - copied_bytes ? SIZE_MAX : copied_bytes + size;
@ -206,6 +212,6 @@ namespace sd {
LOG_DEBUG("%s graph cut cache added %6.2f MB (%zu tensors)", LOG_DEBUG("%s graph cut cache added %6.2f MB (%zu tensors)",
log_desc, copied_bytes / (1024.f * 1024.f), copied_count); log_desc, copied_bytes / (1024.f * 1024.f), copied_count);
} }
return true; return GGML_STATUS_SUCCESS;
} }
} }

8
src/core/runner_cache.h
View File

@ -20,7 +20,8 @@ namespace sd {
~CachedTensor(); ~CachedTensor();
static std::unique_ptr<CachedTensor> copy(ggml_backend_t backend, static std::unique_ptr<CachedTensor> copy(ggml_backend_t backend,
const std::string& name, const std::string& name,
ggml_tensor* source); ggml_tensor* source,
ggml_status& status);
}; };
using CachedTensors = std::map<std::string, std::unique_ptr<CachedTensor>>; using CachedTensors = std::map<std::string, std::unique_ptr<CachedTensor>>;
@ -41,7 +42,8 @@ namespace sd {
const std::map<std::string, ggml_tensor*>& outputs() const { return outputs_; } const std::map<std::string, ggml_tensor*>& outputs() const { return outputs_; }
size_t pending_bytes(ggml_cgraph* graph) const; size_t pending_bytes(ggml_cgraph* graph) const;
size_t resident_bytes(ggml_backend_dev_t device) const; size_t resident_bytes(ggml_backend_dev_t device) const;
bool capture(ggml_cgraph* graph); bool empty() const { return committed_.empty(); }
ggml_status capture(ggml_cgraph* graph);
void graph_end(bool success); void graph_end(bool success);
void clear(); void clear();
}; };
@ -57,7 +59,7 @@ namespace sd {
size_t resident_bytes(ggml_backend_dev_t device) const; size_t resident_bytes(ggml_backend_dev_t device) const;
size_t estimate_output_bytes(ggml_cgraph* graph, size_t estimate_output_bytes(ggml_cgraph* graph,
const ggml_graph_cut::Segment& segment) const; const ggml_graph_cut::Segment& segment) const;
bool capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc); ggml_status capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc);
void prune(const std::unordered_set<std::string>& keep_names); void prune(const std::unordered_set<std::string>& keep_names);
void clear() { tensors_.clear(); } void clear() { tensors_.clear(); }
}; };

2
src/model/diffusion/model.hpp
View File

@ -71,6 +71,8 @@ struct AnimaDiffusionExtra {
struct QwenImage21DiffusionExtra { struct QwenImage21DiffusionExtra {
const sd::Tensor<int32_t>* image_slots = nullptr; const sd::Tensor<int32_t>* image_slots = nullptr;
// Nonzero IDs identify immutable prefix inputs within one sampling run.
uint64_t prefix_id = 0;
}; };
struct WanDiffusionExtra { struct WanDiffusionExtra {

225
src/model/diffusion/qwen_image_2_1.hpp
View File

@ -121,6 +121,18 @@ namespace Qwen {
} }
}; };
struct QwenImage21PrefixCache {
enum class Mode {
NONE,
STORE,
REUSE
};
Mode mode = Mode::NONE;
std::string name;
std::string cut_group;
int64_t prefix_length = 0;
};
class QwenImage21ZeroCenterRMSNorm : public RMSNorm { class QwenImage21ZeroCenterRMSNorm : public RMSNorm {
public: public:
using RMSNorm::RMSNorm; using RMSNorm::RMSNorm;
@ -160,7 +172,7 @@ namespace Qwen {
} }
} }
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::vector<QwenImage21Segment>& segments, const std::vector<ggml_tensor*>& masks) { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::vector<QwenImage21Segment>& segments, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
int64_t heads = x->ne[0] / dim_head; int64_t heads = x->ne[0] / dim_head;
auto project = [&](const char* name) { auto project = [&](const char* name) {
auto h = std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, x); auto h = std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, x);
@ -173,14 +185,36 @@ namespace Qwen {
k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"])->forward(ctx, k); k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"])->forward(ctx, k);
q = Rope::apply_rope(ctx->ggml_ctx, q, pe); q = Rope::apply_rope(ctx->ggml_ctx, q, pe);
k = Rope::apply_rope(ctx->ggml_ctx, k, pe); k = Rope::apply_rope(ctx->ggml_ctx, k, pe);
if (cache.mode == QwenImage21PrefixCache::Mode::STORE) {
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);
// 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);
ctx->persist_cache_tensor(cache.name + "." + name, copy);
};
persist(k, 1, "k");
persist(v, 2, "v");
}
ggml_tensor* result = nullptr; ggml_tensor* result = nullptr;
for (size_t i = 0; i < segments.size(); ++i) { if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) {
const auto& segment = segments[i]; auto prefix_k = ctx->load_cache_tensor(cache.name + ".k");
auto sq = ggml_ext_slice(ctx->ggml_ctx, q, 1, segment.start, segment.end); auto prefix_v = ctx->load_cache_tensor(cache.name + ".v");
auto sk = ggml_ext_slice(ctx->ggml_ctx, k, 1, 0, segment.end); GGML_ASSERT(prefix_k != nullptr && prefix_v != nullptr);
auto sv = ggml_ext_slice(ctx->ggml_ctx, v, 2, 0, segment.end); k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 1);
auto out = ggml_ext_attention_ext(ctx, sq, sk, sv, heads, masks[i], true, ctx->flash_attn_enabled); v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2);
result = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1); result = ggml_ext_attention_ext(ctx, q, k, v, heads, nullptr, true, ctx->flash_attn_enabled);
} 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);
result = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1);
}
} }
auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]); auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
@ -219,13 +253,14 @@ namespace Qwen {
return ggml_concat(ctx, prefix, target, 1); return ggml_concat(ctx, prefix, target, 1);
} }
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector<ggml_tensor*>& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks) { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector<ggml_tensor*>& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
auto h = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm1"])->forward(ctx, x); const int64_t prefix_length = cache.mode == QwenImage21PrefixCache::Mode::REUSE ? 0 : layout.prefix_length;
h = modulate(ctx->ggml_ctx, h, modulation[0], layout.prefix_length); auto h = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm1"])->forward(ctx, x);
h = std::dynamic_pointer_cast<QwenImage21Attention>(blocks["attn"])->forward(ctx, h, pe, layout.segments, masks); h = modulate(ctx->ggml_ctx, h, modulation[0], prefix_length);
x = ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[1], layout.prefix_length, true)); h = std::dynamic_pointer_cast<QwenImage21Attention>(blocks["attn"])->forward(ctx, h, pe, layout.segments, masks, cache);
h = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm2"])->forward(ctx, x); x = ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[1], prefix_length, true));
h = modulate(ctx->ggml_ctx, h, modulation[2], layout.prefix_length); h = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm2"])->forward(ctx, x);
h = modulate(ctx->ggml_ctx, h, modulation[2], prefix_length);
ggml_tensor* gate; ggml_tensor* gate;
auto fused = blocks.find("img_mlp.gate_up"); auto fused = blocks.find("img_mlp.gate_up");
if (fused != blocks.end()) { if (fused != blocks.end()) {
@ -239,7 +274,7 @@ namespace Qwen {
} }
h = ggml_mul(ctx->ggml_ctx, h, ggml_silu(ctx->ggml_ctx, gate)); h = ggml_mul(ctx->ggml_ctx, h, ggml_silu(ctx->ggml_ctx, gate));
h = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.out"])->forward(ctx, h); h = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.out"])->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[3], layout.prefix_length, true)); return ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[3], prefix_length, true));
} }
}; };
@ -261,7 +296,7 @@ namespace Qwen {
} }
} }
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, const std::vector<ggml_tensor*>& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks) { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, const std::vector<ggml_tensor*>& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
auto time = ggml_concat(ctx->ggml_ctx, timestep, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), 0); auto time = ggml_concat(ctx->ggml_ctx, timestep, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), 0);
// Runtime flow timesteps already use the [0, 1000] scale. // Runtime flow timesteps already use the [0, 1000] scale.
time = ggml_ext_timestep_embedding(ctx->ggml_ctx, time, 256, 10000, 1.f); time = ggml_ext_timestep_embedding(ctx->ggml_ctx, time, 256, 10000, 1.f);
@ -269,27 +304,37 @@ namespace Qwen {
time = ggml_silu(ctx->ggml_ctx, time); time = ggml_silu(ctx->ggml_ctx, time);
auto modulation = std::dynamic_pointer_cast<Linear>(blocks["modulation.1"])->forward(ctx, time); auto modulation = std::dynamic_pointer_cast<Linear>(blocks["modulation.1"])->forward(ctx, time);
auto mod = ggml_ext_chunk(ctx->ggml_ctx, modulation, 4, 0); auto mod = ggml_ext_chunk(ctx->ggml_ctx, modulation, 4, 0);
auto text = std::dynamic_pointer_cast<QwenImage21TextProjection>(blocks["txt_in"])->forward(ctx, context);
auto img_in = std::dynamic_pointer_cast<Linear>(blocks["img_in"]); auto img_in = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
ggml_tensor* joint = nullptr; ggml_tensor* joint = nullptr;
for (const auto& segment : layout.segments) { if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) {
ggml_tensor* h; joint = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, x, 1, 1));
if (segment.image_index < 0) { } else {
h = ggml_ext_slice(ctx->ggml_ctx, text, 1, segment.context_start, auto text = std::dynamic_pointer_cast<QwenImage21TextProjection>(blocks["txt_in"])->forward(ctx, context);
segment.context_start + segment.end - segment.start); for (const auto& segment : layout.segments) {
} else { ggml_tensor* h;
auto image = segment.image_index == static_cast<int>(refs.size()) ? x : refs[segment.image_index]; if (segment.image_index < 0) {
h = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, image, 1, 1)); h = ggml_ext_slice(ctx->ggml_ctx, text, 1, segment.context_start,
segment.context_start + segment.end - segment.start);
} else {
auto image = segment.image_index == static_cast<int>(refs.size()) ? x : refs[segment.image_index];
h = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, image, 1, 1));
}
joint = joint == nullptr ? h : ggml_concat(ctx->ggml_ctx, joint, h, 1);
} }
joint = joint == nullptr ? h : ggml_concat(ctx->ggml_ctx, joint, h, 1);
} }
sd::ggml_graph_cut::mark_graph_cut(joint, "qwen_image_2_1.prelude", "joint"); sd::ggml_graph_cut::mark_graph_cut(joint, "qwen_image_2_1.prelude", "joint");
for (int i = 0; i < config.num_layers; ++i) { for (int i = 0; i < config.num_layers; ++i) {
auto block = std::dynamic_pointer_cast<QwenImage21TransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]); const std::string layer = "transformer_blocks." + std::to_string(i);
joint = block->forward(ctx, joint, mod, pe, layout, masks); auto layer_cache = cache;
sd::ggml_graph_cut::mark_graph_cut(joint, "qwen_image_2_1.transformer_blocks." + std::to_string(i), "joint"); layer_cache.name = cache.name + "." + std::to_string(i);
layer_cache.cut_group = "qwen_image_2_1." + layer;
auto block = std::dynamic_pointer_cast<QwenImage21TransformerBlock>(blocks[layer]);
joint = block->forward(ctx, joint, mod, pe, layout, masks, layer_cache);
sd::ggml_graph_cut::mark_graph_cut(joint, layer_cache.cut_group, "joint");
}
if (cache.mode != QwenImage21PrefixCache::Mode::REUSE) {
joint = ggml_ext_slice(ctx->ggml_ctx, joint, 1, layout.prefix_length, joint->ne[1]);
} }
joint = ggml_ext_slice(ctx->ggml_ctx, joint, 1, layout.prefix_length, joint->ne[1]);
auto scale = std::dynamic_pointer_cast<Linear>(blocks["norm_out.linear"])->forward(ctx, ggml_ext_chunk(ctx->ggml_ctx, time, 2, 1)[0]); auto scale = std::dynamic_pointer_cast<Linear>(blocks["norm_out.linear"])->forward(ctx, ggml_ext_chunk(ctx->ggml_ctx, time, 2, 1)[0]);
joint = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out.norm"])->forward(ctx, joint); joint = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out.norm"])->forward(ctx, joint);
joint = ggml_mul(ctx->ggml_ctx, joint, ggml_scale_bias(ctx->ggml_ctx, scale, 1.f, 1.f)); joint = ggml_mul(ctx->ggml_ctx, joint, ggml_scale_bias(ctx->ggml_ctx, scale, 1.f, 1.f));
@ -303,11 +348,18 @@ namespace Qwen {
QwenImage21Model model; QwenImage21Model model;
std::vector<float> pe_data; std::vector<float> pe_data;
std::vector<sd::Tensor<float>> mask_data; std::vector<sd::Tensor<float>> mask_data;
bool prefix_cache_enabled = true;
bool prefix_cache_disabled = false;
QwenImage21Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr) 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), : DiffusionModelRunner(backend, prefix, weight_manager),
config(QwenImage21Config::detect_from_weights(weights, prefix)), config(QwenImage21Config::detect_from_weights(weights, prefix)),
model(config) { model(config) {
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());
}
}
model.init(params_ctx, weights, prefix); model.init(params_ctx, weights, prefix);
} }
@ -317,6 +369,22 @@ namespace Qwen {
model.get_param_tensors(tensors, prefix); model.get_param_tensors(tensors, prefix);
} }
bool has_prefix_cache(const QwenImage21PrefixCache& cache) {
for (int i = 0; i < config.num_layers; ++i) {
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) {
return false;
}
}
return true;
}
sd::Tensor<float> compute(int n_threads, const DiffusionParams& inputs) override { sd::Tensor<float> compute(int n_threads, const DiffusionParams& inputs) override {
const auto& x = tensor_or_empty(inputs.x); const auto& x = tensor_or_empty(inputs.x);
const auto& context = tensor_or_empty(inputs.context); const auto& context = tensor_or_empty(inputs.context);
@ -345,38 +413,75 @@ namespace Qwen {
LOG_ERROR("%s", error.what()); LOG_ERROR("%s", error.what());
return {}; return {};
} }
pe_data = Rope::embed_nd(layout.positions, 1, 10000.f, config.axes_dim); if (!runner_started()) {
mask_data.clear(); prefix_cache_disabled = false;
for (const auto& segment : layout.segments) { }
sd::Tensor<float> mask; QwenImage21PrefixCache cache;
if (segment.image_index < 0) { if (prefix_cache_enabled && !prefix_cache_disabled && extra != nullptr && extra->prefix_id != 0 && layout.prefix_length > 0) {
mask = sd::Tensor<float>::zeros({segment.end, segment.end - segment.start}); cache.name = "qwen_image_2_1.prefix." + std::to_string(extra->prefix_id);
for (int64_t q = segment.start; q < segment.end; ++q) { cache.prefix_length = layout.prefix_length;
for (int64_t k = q + 1; k < segment.end; ++k) { cache.mode = has_prefix_cache(cache) ? QwenImage21PrefixCache::Mode::REUSE : QwenImage21PrefixCache::Mode::STORE;
mask[k + segment.end * (q - segment.start)] = -INFINITY; }
auto run = [&](const QwenImage21PrefixCache& active_cache) {
const bool cached = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE;
const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0);
pe_data = Rope::embed_nd(std::vector<std::vector<float>>(first_position, layout.positions.end()), 1, 10000.f, config.axes_dim);
mask_data.clear();
if (!cached) {
for (const auto& segment : layout.segments) {
sd::Tensor<float> mask;
if (segment.image_index < 0) {
mask = sd::Tensor<float>::zeros({segment.end, segment.end - segment.start});
for (int64_t q = segment.start; q < segment.end; ++q) {
for (int64_t k = q + 1; k < segment.end; ++k) {
mask[k + segment.end * (q - segment.start)] = -INFINITY;
}
}
} }
mask_data.push_back(std::move(mask));
} }
} }
mask_data.push_back(std::move(mask)); auto build = [&]() {
} auto graph = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE * 2);
auto build = [&]() { auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2,
auto graph = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE * 2); layout.positions.size() - (cached ? layout.prefix_length : 0));
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, layout.positions.size()); set_backend_tensor_data(pe, pe_data.data());
set_backend_tensor_data(pe, pe_data.data()); std::vector<ggml_tensor*> masks, ref_inputs;
std::vector<ggml_tensor*> masks, ref_inputs; for (const auto& mask : mask_data) {
for (const auto& mask : mask_data) { masks.push_back(mask.empty() ? nullptr : make_input(mask));
masks.push_back(mask.empty() ? nullptr : make_input(mask)); }
} if (!cached) {
for (const auto& ref : refs) { for (const auto& ref : refs) {
ref_inputs.push_back(make_input(ref)); ref_inputs.push_back(make_input(ref));
} }
auto ctx = get_context(); }
auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), make_input(context), auto ctx = get_context();
ref_inputs, pe, layout, masks); auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), cached ? nullptr : make_input(context),
ggml_build_forward_expand(graph, out); ref_inputs, pe, layout, masks, active_cache);
return graph; ggml_build_forward_expand(graph, out);
return graph;
};
return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), x.dim());
}; };
return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), x.dim()); auto 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.
free_cache_ctx_and_buffer();
prefix_cache_disabled = true;
LOG_WARN("Qwen Image 2.1: insufficient memory for prefix caching; retrying without it for this sampling run");
return run(QwenImage21PrefixCache{});
}
if (!result.empty() && cache.mode == QwenImage21PrefixCache::Mode::STORE) {
if (!has_prefix_cache(cache)) {
free_cache_ctx_and_buffer();
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);
}
}
return result;
} }
}; };
} }

22
src/pipeline/diffusion_engine.cpp
View File

@ -7,6 +7,7 @@
#include <list> #include <list>
#include <mutex> #include <mutex>
#include <set> #include <set>
#include <tuple>
#include <type_traits> #include <type_traits>
#include <unordered_set> #include <unordered_set>
#include <utility> #include <utility>
@ -2255,6 +2256,15 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
}; };
RunnerEndOnExit sample_diffusion_runner_end{work_diffusion_model.get()}; RunnerEndOnExit sample_diffusion_runner_end{work_diffusion_model.get()};
// These inputs are immutable for this sampling run. Extensions may replace or
// modify them per step, so those paths need an explicit stability contract first.
const bool cache_qwen_prefix = version == VERSION_QWEN_IMAGE_2_1 &&
std::none_of(generation_extensions.begin(), generation_extensions.end(),
[](const auto& extension) { return extension->is_enabled(); });
using QwenPrefixInputs = std::tuple<const sd::Tensor<float>*, const sd::Tensor<int32_t>*,
const std::vector<sd::Tensor<float>>*>;
std::vector<QwenPrefixInputs> qwen_prefix_inputs;
RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr}; RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr};
const bool apply_denoise_mask = !denoise_mask.empty() && const bool apply_denoise_mask = !denoise_mask.empty() &&
@ -2524,6 +2534,18 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
extension->before_diffusion(diffusion_params, step); extension->before_diffusion(diffusion_params, step);
} }
if (cache_qwen_prefix) {
auto* extra = std::get_if<QwenImage21DiffusionExtra>(&diffusion_params.extra);
if (extra != nullptr) {
auto key = std::make_tuple(diffusion_params.context, extra->image_slots,
diffusion_params.ref_image_params.pass_to_dit ? diffusion_params.ref_latents : nullptr);
auto entry = std::find(qwen_prefix_inputs.begin(), qwen_prefix_inputs.end(), key);
extra->prefix_id = static_cast<uint64_t>(entry - qwen_prefix_inputs.begin()) + 1;
if (entry == qwen_prefix_inputs.end()) {
qwen_prefix_inputs.push_back(key);
}
}
}
auto output_opt = work_diffusion_model->compute(n_threads, diffusion_params); auto output_opt = work_diffusion_model->compute(n_threads, diffusion_params);
if (output_opt.empty()) { if (output_opt.empty()) {
LOG_ERROR("diffusion model compute failed"); LOG_ERROR("diffusion model compute failed");

3
src/pipeline/model_builders.cpp
View File

@ -291,7 +291,8 @@ namespace sd::model_builders {
result.diffusion = std::make_shared<Qwen::QwenImage21Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), result.diffusion = std::make_shared<Qwen::QwenImage21Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map, tensor_storage_map,
"model.diffusion_model", "model.diffusion_model",
weight_manager); weight_manager,
sd_ctx_params->model_args);
} else { } else {
result.diffusion = std::make_shared<Qwen::QwenImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), result.diffusion = std::make_shared<Qwen::QwenImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map, tensor_storage_map,