diff --git a/docs/qwen_image_2.1.md b/docs/qwen_image_2.1.md
index 3eaae207..56999bcc 100644
--- a/docs/qwen_image_2.1.md
+++ b/docs/qwen_image_2.1.md
@@ -23,7 +23,7 @@ Run the following commands from the build directory. Use image dimensions divisi
### Text to image
```powershell
-.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\qwen_image_2.1_int8_convrot.safetensors --vae ..\models\vae\qwen_image_2.1_vae_bf16.safetensors --llm ..\models\text_encoders\Qwen3VL-8B-Instruct-Q4_K_M.gguf -p "a lovely cat holding a sign says 'qwen2.1.cpp'" --cfg-scale 6.0 --sampling-method euler -v --offload-to-cpu -o qwen_image_2.1.png
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\qwen_image_2.1_int8_convrot.safetensors --vae ..\models\vae\qwen_image_2.1_vae_bf16.safetensors --llm ..\models\text_encoders\Qwen3VL-8B-Instruct-Q4_K_M.gguf -p "a lovely cat holding a sign says 'qwen2.1.cpp'" --cfg-scale 6.0 --sampling-method euler -v --offload-to-cpu --fa -o qwen_image_2.1.png
```
@@ -35,7 +35,7 @@ To use GGUF diffusion weights, set `--diffusion-model` to the path of a file suc
Pass the reference image with `-r` and describe the edit in `-p`. Vision weights are required; the example below loads them separately with `--llm_vision`.
```powershell
-.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\qwen_image_2.1_int8_convrot.safetensors --vae ..\models\vae\qwen_image_2.1_vae_bf16.safetensors --llm ..\models\text_encoders\Qwen3VL-8B-Instruct-Q4_K_M.gguf --llm_vision ..\models\text_encoders\Qwen3VL-8B-Instruct-mmproj-BF16.gguf -r ..\assets\qwen\qwen_image_2.1.png -p "change 'qwen2.1.cpp' to 'sd.cpp'" --cfg-scale 6.0 --sampling-method euler -v --offload-to-cpu -o qwen_image_2.1_edit.png
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\qwen_image_2.1_int8_convrot.safetensors --vae ..\models\vae\qwen_image_2.1_vae_bf16.safetensors --llm ..\models\text_encoders\Qwen3VL-8B-Instruct-Q4_K_M.gguf --llm_vision ..\models\text_encoders\Qwen3VL-8B-Instruct-mmproj-BF16.gguf -r ..\assets\qwen\qwen_image_2.1.png -p "change 'qwen2.1.cpp' to 'sd.cpp'" --cfg-scale 6.0 --sampling-method euler -v --offload-to-cpu --fa -o qwen_image_2.1_edit.png
```
For multiple reference images, repeat `-r` in the desired order, for example `-r first.png -r second.png`.
@@ -44,7 +44,27 @@ For multiple reference images, repeat `-r` in the desired order, for example `-r
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.
+Set `qwen_image_2_1_prefix_cache_type` in `--model-args` to `auto` or a type name using the same parser and case-sensitive names as `--type`:
+
+- `auto` (default): use FP16 only when Flash Attention is enabled, Sage Attention is disabled, the attention scale is unchanged, and every attention operation in the cache-writing or cache-reading graph selects Flash Attention after backend support checks. If an operation falls back, rebuild the prefix in FP32 before executing and keep FP32 for the rest of that sampling run.
+- `f32`: always store FP32 keys and values.
+- `f16`: always store FP16 keys and values, including with ordinary attention or custom attention scaling. This saves cache memory but can introduce additional rounding error.
+- Other types, such as `bf16`, `q4_1`, `q5_0`, `q5_1`, `q8_0`, `q4_K`, `q6_K`, `iq4_nl`, and `iq4_xs`: use the requested storage type if the ggml build provides runtime conversion to and from FP32. Quantization is lossy and must be selected explicitly; `auto` never selects a quantized type.
+
+Cache data is packed into contiguous rows of `hidden_size` elements before conversion, so 256-element quantization blocks work with the model's 128-element attention heads without padding. The type's block size must divide `hidden_size`. Unknown types, types lacking runtime conversion (for example `q8_1` and several IQ formats), and incompatible block sizes are ignored with a warning, leaving the previous setting or the default `auto` unchanged.
+
+For example, use `--model-args qwen_image_2_1_prefix_cache_type=q8_0` to enable 8-bit cache storage. Cached keys and values are converted back to the attention input type before concatenating with the current target tokens. This reduces persistent cache memory; attention working buffers still use floating-point values, and conversion adds work on each step. Backends without the required conversion operations use the existing CPU fallback.
+
+For the default 32-layer model, a prefix of 4096 tokens takes approximately the following memory per condition, excluding weights, working buffers, and allocation overhead:
+
+| Cache type | Memory |
+| --- | ---: |
+| `f32` | 4 GiB |
+| `f16` | 2 GiB |
+| `q8_0` | 1.0625 GiB |
+| `q4_0` | 0.5625 GiB |
+
+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.
diff --git a/examples/common/common.cpp b/examples/common/common.cpp
index 24dc937b..fa7354bc 100644
--- a/examples/common/common.cpp
+++ b/examples/common/common.cpp
@@ -518,7 +518,8 @@ ArgOptions SDContextParams::get_options() {
{"",
"--model-args",
"extra model args, key=value list. Supports chroma_use_dit_mask, chroma_use_t5_mask, "
- "chroma_t5_mask_pad, qwen_image_zero_cond_t, qwen_image_2_1_prefix_cache",
+ "chroma_t5_mask_pad, qwen_image_zero_cond_t, qwen_image_2_1_prefix_cache, "
+ "qwen_image_2_1_prefix_cache_type (auto or a type name from --type)",
(int)',',
&model_args},
{"",
diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp
index 74324169..c4c88a63 100644
--- a/src/core/ggml_extend.cpp
+++ b/src/core/ggml_extend.cpp
@@ -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,
diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h
index 33d8c3bf..be88c58a 100644
--- a/src/core/ggml_extend.h
+++ b/src/core/ggml_extend.h
@@ -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,
diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp
index 59573210..775d9420 100644
--- a/src/core/ggml_runner.cpp
+++ b/src/core/ggml_runner.cpp
@@ -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> GGMLRunner::compute(get_graph_cb_t get_graph,
diff --git a/src/core/ggml_runner.h b/src/core/ggml_runner.h
index 9084db10..02630c0d 100644
--- a/src/core/ggml_runner.h
+++ b/src/core/ggml_runner.h
@@ -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);
diff --git a/src/model/diffusion/control.hpp b/src/model/diffusion/control.hpp
index 34387d1b..3b08ed15 100644
--- a/src/model/diffusion/control.hpp
+++ b/src/model/diffusion/control.hpp
@@ -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);
diff --git a/src/model/diffusion/qwen_image_2_1.hpp b/src/model/diffusion/qwen_image_2_1.hpp
index 9884325c..09311c29 100644
--- a/src/model/diffusion/qwen_image_2_1.hpp
+++ b/src/model/diffusion/qwen_image_2_1.hpp
@@ -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 pe_data;
std::vector> 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_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 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 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(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;
diff --git a/src/model/diffusion/sensenova_u1.h b/src/model/diffusion/sensenova_u1.h
index 22b789fa..1702ed85 100644
--- a/src/model/diffusion/sensenova_u1.h
+++ b/src/model/diffusion/sensenova_u1.h
@@ -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,
diff --git a/src/model/vae/ltx_vae.hpp b/src/model/vae/ltx_vae.hpp
index ac22c09b..f680f200 100644
--- a/src/model/vae/ltx_vae.hpp
+++ b/src/model/vae/ltx_vae.hpp
@@ -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(feat_map.size()); ++feat_idx) {
ggml_tensor* feat_cache = feat_map[static_cast(feat_idx)];
if (feat_cache != nullptr) {
- cache(temporal_feat_cache_name(static_cast(feat_idx)), feat_cache);
- ggml_build_forward_expand(gf, feat_cache);
+ runner_ctx.persist_cache_tensor(temporal_feat_cache_name(static_cast(feat_idx)), feat_cache);
}
}
diff --git a/src/model/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp
index 8d14cf95..73cd0953 100644
--- a/src/model/vae/wan_vae.hpp
+++ b/src/model/vae/wan_vae.hpp
@@ -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);
}
}