CUDA: on 12 GB GPUs the final VAE decode falls back to the CPU at 1024x1024

Start with iris_decode_latent() in iris.c and the VAE device allocation and decode paths in iris_vae.c; trace how vae_decode_dev() handles allocation failure and how GPU memory is reserved and released. Reproduce the reported 1024×1024 CUDA decode on a 12 GB GPU, then check that it decodes on the GPU without slowing generations when the decoder already fits. The report includes make test and run_test.py --full results to compare.

On an RTX 3060 12 GB with the CUDA build of ff97831, FLUX.2-klein-4B at 1024x1024 decodes the image on the CPU, and that is most of the time:

$ ./iris -d flux-klein-4b -p "a small fishing boat on white background, watercolor" \
    -W 1024 -H 1024 -S 42 -v -o boat.png
...
  Total denoising: 5726.2 ms (5.73 s)
Decoding image...VAE: GPU decode failed, using the CPU
 done (134.5s)
Generated in 143.8s total

After denoising the transformer is still on the GPU (7.2 GiB of bf16, the process peaks at 8209 MiB), and at 1024x1024 the GPU decoder needs about 5.4 GiB more: four 1 GiB activation buffers (the largest tensor is 1024x1024x256 floats), 1 GiB of attention scratch, 256 MiB of im2col scratch. vae_decode_dev() fails and iris_vae_decode() goes to the CPU. 768x1152, img2img and the interactive CLI do the same.

Freeing the transformer before the decode fixes it (1.2 s on the GPU), but doing it always would reload the transformer at every CLI image also on GPUs where everything fits. It is also not enough in the CLI: if the transformer is freed before every decode, the decoder buffers stay allocated and the second image gets "FLUX.2: GPU weight upload failed, using the CPU".

The patch below does it only when the decoder does not fit:

• iris_vae_dev_reserve() allocates the GPU decoder (weights and activation buffers) for the latent size, with the same vdev_load() and vdev_ensure_buffers() the decoder uses, and returns 0 if they do not fit. Without a GPU decoder (CPU and MPS builds, no GPU) it returns 1.
• iris_decode_latent(), already public and now used by all the final decodes in iris.c, calls it first. On failure it frees the transformer, which the next generation loads again through the existing *_transformer_if_needed() functions, decodes on the GPU, then frees the decoder activation buffers (iris_vae_dev_release()) so that the next text encoder and transformer find the GPU as for the first image.

When the decoder fits nothing changes: the reserve does the allocations the decoder would do anyway, in the same order, and nothing is freed, so the transformer stays loaded across generations. The step previews in iris_sample.c still call iris_vae_decode(), so the transformer is never freed while sampling.

Numbers: RTX 3060 12 GB, i5-10600 (6 CPUs, shared, so the CPU decode times are high), CUDA 12.6.3, make cuda GPU_ARCH=sm_86 OPENBLAS=1, -S 42. Total is "Total generation time" (CLI: per image). VRAM is nvidia-smi every 0.2 s, minus 279 MiB of other idle processes.

For img2img the reference is resized to 720x720 and encoded on the GPU in both builds.

Images: patched vs ff97831 differ by at most 2, mean 0.04-0.09 (GPU vs CPU decode of the same latent). The 512x512 images, one-shot and CLI, are bit-identical with and without the patch. In the 1024x1024 CLI run the second ff97831 image also had the text encoder on the CPU ("Qwen3: GPU allocation failed, using the CPU"), so it differs from the same prompt generated alone (mean 9.4); the patched second image is bit-identical to it.

In the patched CLI the second image reloads the transformer (1.4 s from the page cache): on a 12 GB GPU the transformer and the decoder take turns. Between images the process goes 8083 MiB (denoising) -> 5699 MiB (decode) -> 1539 MiB (attention and im2col scratch, kept as now).

Tests: make test 3/3 with the patched CUDA build, same mean_diff/max_diff as ff97831 (4.11/60, 6.99/43, 8.17/206); run_test.py --full 4/4 (its 1024x1024 runs take the new path, car checked by eye); the BLAS build passes the same 3 tests (here only with a timeout longer than 300 s, the CPU being shared). make generic, make blas and make cuda: the same 10 warnings as ff97831.

Not tested: MPS (VAE_DEV is not defined there, so the new functions do nothing; I only ran gcc -fsyntax-only -DUSE_METAL -DIRIS_DEV on iris.c and iris_vae.c); ROCm (the changed C files compile as for CUDA); FLUX.2-9B and FLUX.3.

Two things seen while testing, the same with and without the patch:

1. At 896x896 the decoder buffers fit next to the transformer but the 900 MiB attention scratch does not: iris_dev_attention() prints "iris_dev: attention scratch allocation failed" and returns, and the decode goes on with the mid-block attention output not computed. The image looks right but differs from a full GPU decode (max 66, mean 2.7). iris_dev_conv2d() and iris_dev_group_norm() return the same way when their scratch can't be allocated.
2. In the CLI the second image's text encoder runs on the CPU (+27 s at 512x512 here): the transformer stays resident, as it should, and Qwen3 does not fit next to it in 12 GB.

The patch applies with git apply on ff97831; glad to open a PR if you prefer.

diff --git a/iris.c b/iris.c
index 721b7e0..b4d2aac 100644
--- a/iris.c
+++ b/iris.c
@@ -45,6 +45,8 @@ extern float *iris_vae_encode(iris_vae_t *vae, const float *img,
                               int batch, int H, int W, int *out_h, int *out_w);
 extern iris_image *iris_vae_decode(iris_vae_t *vae, const float *latent,
                                    int batch, int latent_h, int latent_w);
+extern int iris_vae_dev_reserve(iris_vae_t *vae, int latent_h, int latent_w);
+extern void iris_vae_dev_release(iris_vae_t *vae);
 extern float *iris_image_to_tensor(const iris_image *img);
 
 extern iris_transformer_flux_t *iris_transformer_load_flux(FILE *f);
@@ -535,15 +537,23 @@ iris_ctx *iris_load_dir(const char *model_dir) {
     return ctx;
 }
 
+/* Free the transformer, loaded again on demand by the next generation. */
+static void iris_free_transformers(iris_ctx *ctx) {
+    iris_transformer_free_flux(ctx->transformer);
+    iris_transformer_free_zimage(ctx->zi_transformer);
+    iris_transformer_free_flux3(ctx->flux3_transformer);
+    ctx->transformer = NULL;
+    ctx->zi_transformer = NULL;
+    ctx->flux3_transformer = NULL;
+}
+
 void iris_free(iris_ctx *ctx) {
     if (!ctx) return;
 
     iris_tokenizer_free(ctx->tokenizer);
     qwen3_encoder_free(ctx->qwen3_encoder);
     iris_vae_free(ctx->vae);
-    iris_transformer_free_flux(ctx->transformer);
-    iris_transformer_free_zimage(ctx->zi_transformer);
-    iris_transformer_free_flux3(ctx->flux3_transformer);
+    iris_free_transformers(ctx);
 
     free(ctx);
 }
@@ -819,9 +829,7 @@ static iris_image *iris_generate_zimage_with_embeddings(iris_ctx *ctx,
     /* Decode latent to image */
     iris_image *img = NULL;
     if (ctx->vae) {
-        if (iris_phase_callback) iris_phase_callback("decoding image", 0);
-        img = iris_vae_decode(ctx->vae, latent, 1, post_h, post_w);
-        if (iris_phase_callback) iris_phase_callback("decoding image", 1);
+        img = iris_decode_latent(ctx, latent, post_h, post_w);
     }
 
     free(latent);
@@ -890,9 +898,7 @@ static iris_image *iris_generate_flux3_with_embeddings(iris_ctx *ctx,
         return NULL;
     }
 
-    if (iris_phase_callback) iris_phase_callback("decoding image", 0);
-    iris_image *img = iris_vae_decode(ctx->vae, latent, 1, latent_h, latent_w);
-    if (iris_phase_callback) iris_phase_callback("decoding image", 1);
+    iris_image *img = iris_decode_latent(ctx, latent, latent_h, latent_w);
     free(latent);
     return img;
 }
@@ -1072,9 +1078,7 @@ static iris_image *iris_generate_flux3_refs(iris_ctx *ctx, const char *prompt,
         set_error("Sampling failed");
         goto done;
     }
-    if (iris_phase_callback) iris_phase_callback("decoding image", 0);
-    img = iris_vae_decode(ctx->vae, latent, 1, latent_h, latent_w);
-    if (iris_phase_callback) iris_phase_callback("decoding image", 1);
+    img = iris_decode_latent(ctx, latent, latent_h, latent_w);
     free(latent);
 
     /* A single padded reference: crop the output back to the image. */
@@ -1253,9 +1257,7 @@ iris_image *iris_generate(iris_ctx *ctx, const char *prompt,
     /* Decode latent to image */
     iris_image *img = NULL;
     if (ctx->vae) {
-        if (iris_phase_callback) iris_phase_callback("decoding image", 0);
-        img = iris_vae_decode(ctx->vae, latent, 1, latent_h, latent_w);
-        if (iris_phase_callback) iris_phase_callback("decoding image", 1);
+        img = iris_decode_latent(ctx, latent, latent_h, latent_w);
     }
 
     free(latent);
@@ -1356,9 +1358,7 @@ iris_image *iris_generate_with_embeddings(iris_ctx *ctx,
     /* Decode latent to image */
     iris_image *img = NULL;
     if (ctx->vae) {
-        if (iris_phase_callback) iris_phase_callback("decoding image", 0);
-        img = iris_vae_decode(ctx->vae, latent, 1, latent_h, latent_w);
-        if (iris_phase_callback) iris_phase_callback("decoding image", 1);
+        img = iris_decode_latent(ctx, latent, latent_h, latent_w);
     } else {
         set_error("No VAE loaded");
         free(latent);
@@ -1452,9 +1452,7 @@ iris_image *iris_generate_with_embeddings_and_noise(iris_ctx *ctx,
     /* Decode latent to image */
     iris_image *img = NULL;
     if (ctx->vae) {
-        if (iris_phase_callback) iris_phase_callback("decoding image", 0);
-        img = iris_vae_decode(ctx->vae, latent, 1, latent_h, latent_w);
-        if (iris_phase_callback) iris_phase_callback("decoding image", 1);
+        img = iris_decode_latent(ctx, latent, latent_h, latent_w);
     } else {
         set_error("No VAE loaded");
         free(latent);
@@ -1733,9 +1731,7 @@ iris_image *iris_img2img(iris_ctx *ctx, const char *prompt,
     /* Decode */
     iris_image *result = NULL;
     if (ctx->vae) {
-        if (iris_phase_callback) iris_phase_callback("decoding image", 0);
-        result = iris_vae_decode(ctx->vae, latent, 1, out_lat_h, out_lat_w);
-        if (iris_phase_callback) iris_phase_callback("decoding image", 1);
+        result = iris_decode_latent(ctx, latent, out_lat_h, out_lat_w);
     }
 
     free(latent);
@@ -1969,9 +1965,7 @@ iris_image *iris_multiref(iris_ctx *ctx, const char *prompt,
     /* Decode */
     iris_image *result = NULL;
     if (ctx->vae) {
-        if (iris_phase_callback) iris_phase_callback("decoding image", 0);
-        result = iris_vae_decode(ctx->vae, latent, 1, latent_h, latent_w);
-        if (iris_phase_callback) iris_phase_callback("decoding image", 1);
+        result = iris_decode_latent(ctx, latent, latent_h, latent_w);
     }
 
     free(latent);
@@ -2052,7 +2046,15 @@ iris_image *iris_decode_latent(iris_ctx *ctx, const float *latent,
                                int latent_h, int latent_w) {
     if (!ctx || !latent || !ctx->vae) return NULL;
     if (iris_phase_callback) iris_phase_callback("decoding image", 0);
+    /* The GPU decoder needs about 5 GB at 1024x1024. If that does not fit
+     * next to the transformer weights, free the transformer (the next
+     * generation loads it again) instead of falling back to the much
+     * slower CPU decoder, then free the decoder buffers to leave room for
+     * the next text encoder and transformer. */
+    int no_room = !iris_vae_dev_reserve(ctx->vae, latent_h, latent_w);
+    if (no_room) iris_free_transformers(ctx);
     iris_image *img = iris_vae_decode(ctx->vae, latent, 1, latent_h, latent_w);
+    if (no_room) iris_vae_dev_release(ctx->vae);
     if (iris_phase_callback) iris_phase_callback("decoding image", 1);
     return img;
 }
@@ -2173,9 +2175,7 @@ iris_image *iris_img2img_debug_py(iris_ctx *ctx, const iris_params *params) {
     /* Decode */
     iris_image *result = NULL;
     if (ctx->vae) {
-        if (iris_phase_callback) iris_phase_callback("decoding image", 0);
-        result = iris_vae_decode(ctx->vae, latent, 1, latent_h, latent_w);
-        if (iris_phase_callback) iris_phase_callback("decoding image", 1);
+        result = iris_decode_latent(ctx, latent, latent_h, latent_w);
     }
 
     free(latent);
diff --git a/iris_vae.c b/iris_vae.c
index c540cf9..7d97631 100644
--- a/iris_vae.c
+++ b/iris_vae.c
@@ -1325,6 +1325,29 @@ static float *vae_encode_dev(iris_vae_t *vae, const float *img, int H, int W,
 }
 #endif /* VAE_DEV */
 
+/* Allocate the GPU decoder, weights and activation buffers, for a
+ * latent_h x latent_w latent. Returns 0 if the GPU memory is not enough,
+ * 1 otherwise (also when there is no GPU decoder). */
+int iris_vae_dev_reserve(iris_vae_t *vae, int latent_h, int latent_w) {
+#ifdef VAE_DEV
+    if (iris_dev_available())
+        return (vae->dev || vdev_load(vae)) &&
+               vdev_ensure_buffers(vae, latent_h * 2, latent_w * 2);
+#else
+    (void)vae; (void)latent_h; (void)latent_w;
+#endif
+    return 1;
+}
+
+/* Free the GPU decoder activation buffers, allocated again on next use. */
+void iris_vae_dev_release(iris_vae_t *vae) {
+#ifdef VAE_DEV
+    if (vae->dev) vdev_free_buffers(vae->dev);
+#else
+    (void)vae;
+#endif
+}
+
 iris_image *iris_vae_decode(iris_vae_t *vae, const float *latent,
                             int batch, int latent_h, int latent_w) {
 #ifdef USE_METAL

Disclosure: written by an AI agent (Claude), which ran the tests above; posted from this account with its owner's permission.