Qwen Image 2.1 prefix cache: store K/V as F16 when flash attention is on (half the VRAM, same output)

[brt16]

Thanks for implementing the prefix cache so quickly in #2035 (from #2031). It works great.

One suggestion: the cache always stores K/V as FP32 (qwen_image_2_1.hpp#L194). But when flash attention is enabled, ggml_ext_attention_ext converts K and V to F16 right before the kernel (ggml_extend.cpp#L675 and #L685). So with flash attention the stored FP32 values are rounded to F16 on every step anyway. Storing them as F16 in the first place would give identical results with half the memory.

The same holds only when:

• flash attention is on
• sage attention is off (its path uses F32 K)
• kv_scale == 1, since the scale is applied before the cast (#L673)

In every other case FP32 is still needed.

Why it matters: the docs estimate about 4 GiB per condition for a 4096-token prefix. Multi-reference edits reach that quickly. Two 1024x1024 references are over 8192 tokens, which is about 8+ GiB at FP32 versus 4+ GiB at F16. On a 16 GB card that difference decides whether the cache fits next to the weights.

The patch attached to #2031 did this (prefix_cache_type()). On Vulkan (RX 6800 XT, --diffusion-fa), its output was pixel-identical to the uncached path. I haven't tested other backends.

[stduhpf]

Having an option for storing the KV cache in a quantized type even with full precision attention might be worth considering to save on VRAM, it's lossy, but might be good enough for it not to matter that much.

[Green-Sky]

+1 for kv quantization. I tried this model too and it came to a crawl with larger image gen, where it has to stream much of the model due to larger "compute buffers".
Also on llama.cpp we have q8_0 with hadamard rotations, which is almost as good as f16, which could make a usable big difference for sd.cpp too. But this would need testing first, as its not the same.

[leejet]

When using FP16 with FA, fallback cases need some additional handling, so I didn’t address them in the initial version.

I’ve now added that handling, along with support for user-specified cache types, and also fixed the issue where the compute buffer size became too large when caching was enabled.