Z-Image on 12 GB GPUs: stream the layers that do not fit instead of running everything on the CPU

The proposed changes are in iris_transformer_zimage.c and README.md; start by reading the Z-Image device loader and forward path around zi_dev_load_block() and zi_forward_dev(). Review the supplied patch and its stated make test results, then check that streamed layers preserve output and handle activation allocation failures. Done means partial GPU offload works as described while existing full-GPU behavior remains unchanged.

On a 12 GB card the Z-Image-Turbo transformer does not fit, and iris runs all of it on the CPU: about 19 minutes per step at 1024x1024. The patch below keeps the layers that fit on the GPU and uploads the others at each step. With it, and with #54 for the final decode, the same card makes a 1024x1024 image in 42 s total, 2.6 s per step.

Setup: RTX 3060 12 GB (sm_86, PCIe 3.0 x16), i5-10600, 32 GB RAM, CUDA 12.6.3, make cuda GPU_ARCH=sm_86 OPENBLAS=1 at ff97831. Model: Tongyi-MAI/Z-Image-Turbo (official F32 safetensors). Run: -W 1024 -H 1024 -S 42 -v, default 9 steps.

What happens now (ff97831):

Encoding text... done (2.0s)
Loading Z-Image transformer...  Loading Z-Image transformer (3 shards)...
  Z-Image: CPU mmap mode enabled (zero-copy f32 weights)
Z-Image: GPU weight upload failed, using the CPU
  Z-Image transformer loaded: dim=3840, heads=30, layers=2+2+30, ffn=10240
 done (10.3s)

Qwen3 runs on the GPU, then one step takes 1157 s on the CPU, about 36 s for each noise refiner and main block (9 steps: about 2 h 54 min).

Why:

• The checkpoint is F32: 6.15B parameters, 24.6 GB. zi_dev_load() uploads a bf16 copy: 11738 MiB of matrices plus 4 MB of f32 vectors.
• cudaMemGetInfo() reports 11911 MiB in total on the 3060 (12288 minus 377 reserved by the driver), a CUDA context with cuBLAS takes about 120 MiB, and the activations at 1024x1024 need another 0.4 GB (zi_dev_ensure() allocates about 98 KiB per token, 4096 image tokens plus the caption). So weights and activations do not fit even on an empty card; here, with 279 MiB used by other idle processes, not even the weights fit.
• zi_dev_load() is all or nothing: a failed upload frees everything and leaves tf->dev NULL, so iris_transformer_forward_zimage() takes the CPU path at every step. The FLUX.2 and FLUX.3 loaders do the same.
• Running the layers that do not fit on the CPU would cost about 36 s per layer per step; uploading one costs about 0.03 s.

The change (iris_transformer_zimage.c and a README note):

• When the GPU is full, the last main layers keep their four big matrices (qkv, proj, ff_in, ff_out: 354 MB per layer) in host memory as bf16, and upload them into a single slot buffer right before they run. Their norms and adaLN weights stay on the GPU; the per-step modulation precompute and zi_dev_block() are unchanged.
• At load, the first layer that does not fit is streamed, and so are all the layers after it.
• In zi_forward_dev(), zi_dev_ensure() runs in a loop: if the activations for the requested size do not fit, zi_dev_evict() streams one more resident layer and retries. A larger image, or another process holding part of the GPU, just means more streamed layers.
• When the number of resident layers changes iris prints one line: Z-Image: 26 of 30 layers on the GPU, the others uploaded at each step. It is printed by the first forward pass, so it ends up after the "Step 1/9" progress text.
• No new API, no option. When everything fits nothing changes: zi_dev_load_block(..., 0) is the same upload as before.
• Streamed layers use the same bf16 rounding and the same kernels, so the output does not depend on the number of streamed layers.

At 1024x1024 the VAE decoder does not fit next to the streamed transformer, so with this patch alone the final decode falls back to the CPU (VAE: GPU decode failed, using the CPU, 133 s here instead of 1.7 s). #54 fixes that for every model by freeing the transformer when the decoder does not fit; the next generation reloads it (about 11 s here).

Same RTX 3060, 1024x1024, 9 steps, seed 42. Total is "Total generation time", load included. Peak VRAM is the whole card, including 279 MiB of other idle processes.

• RSS is mostly the mmapped F32 checkpoint, which stays mapped for the CPU fallback and is the source of the streamed copies. The anonymous memory includes the host copies of the 4 streamed layers (4 x 354 MB).
• 26 of 30 layers stay on the GPU at 1024x1024, 27 at 512x512, where the run takes 23.1 s, 0.66 s per step (ff97831: 148 s for one step).
• Cost of streaming: with 2.5 GB of VRAM held by another process, 18 layers stay resident and a step takes 2.88 s instead of 2.63 s: 32 ms per streamed layer, about 11 GB/s through the pinned staging buffers.
• Determinism: the two runs above and the one with 18 resident layers gave byte-identical PNGs.
• make test (FLUX.2-klein-4B and the Z-Image smoke test): 4/4, same mean_diff/max_diff as ff97831 (4.11/60, 6.99/43, 8.17/206). On ff97831 the Z-Image smoke test runs on the CPU: 82 s for the whole test, 38 s with the patch.
• make generic, make blas, make cuda: the same 10 warnings as ff97831, none in the new code.

Not tested:

• ROCm and Metal: the patch only uses iris_dev_alloc(), iris_dev_upload() and iris_dev_free(), but I could not run those backends.
• GPUs that hold the whole model: the new path is not taken there.
• Interactive mode with growing image sizes: layers are streamed when needed, never moved back to the GPU.
• FLUX.2-9B and FLUX.3 on small GPUs, which have the same all-or-nothing loader.

Related: #37 (Z-Image stuck or killed while loading on 16 GB Macs) is a different code path (the old Metal f16 conversion cache). I found no issue or PR about partial GPU offload.

The patch applies with git apply on ff97831, also on top of #54; glad to open a PR if you prefer.

diff --git a/README.md b/README.md
index 4abe5c1..ca4e221 100644
--- a/README.md
+++ b/README.md
@@ -303,9 +303,12 @@ make rocm       # AMD GPU (needs ROCm, with hipBLAS)
 The CUDA and ROCm backends run the whole pipeline of every model on the
 GPU (text encoder, transformer, VAE), in bf16 with f32 accumulation. Add
 `OPENBLAS=1` to use OpenBLAS for the CPU fallback (used when the GPU runs
-out of memory). The build targets the GPU of the machine it runs on; set
-`GPU_ARCH` (e.g. `GPU_ARCH=sm_121` or `GPU_ARCH=gfx1151`) to build for
-another one, and `CUDA_HOME` / `ROCM_HOME` for non-standard install paths.
+out of memory). On GPUs that cannot hold the Z-Image transformer and its
+activations, such as 12 GB cards, the layers that do not fit stay in host
+memory and are uploaded each time they run. The build targets the GPU of
+the machine it runs on; set `GPU_ARCH` (e.g. `GPU_ARCH=sm_121` or
+`GPU_ARCH=gfx1151`) to build for another one, and `CUDA_HOME` /
+`ROCM_HOME` for non-standard install paths.
 
 For `make blas` on Linux, install OpenBLAS first:
 ```bash
diff --git a/iris_transformer_zimage.c b/iris_transformer_zimage.c
index 3317c1e..5eeaddd 100644
--- a/iris_transformer_zimage.c
+++ b/iris_transformer_zimage.c
@@ -1694,10 +1694,16 @@ typedef struct {
     uint16_t *qkv, *proj, *ff_in, *ff_out, *mod;
     float *norm_q, *norm_k, *attn_norm1, *attn_norm2, *ffn_norm1, *ffn_norm2;
     float *mod_bias;
+    uint16_t *host;             /* streamed layer: qkv, proj, ff_in, ff_out on the host */
 } zi_dev_block_t;
 
+/* When the weights do not fit, the last main layers keep their four big
+ * matrices in host memory and are uploaded to the slot buffer each time they
+ * run: on a 12 GB GPU a few of the 30 layers, 354 MB each. */
 struct zi_dev {
     zi_dev_block_t *noise, *context, *layers;
+    int resident, shown;        /* main layers [0, resident) are on the GPU */
+    uint16_t *slot;
     uint16_t *t_lin1, *t_lin2, *img_in, *cap_in, *mod_final, *proj_final;
     float *t_bias1, *t_bias2, *img_bias, *cap_bias, *cap_norm;
     float *img_pad, *cap_pad, *mod_final_bias, *proj_final_bias;
@@ -1714,9 +1720,19 @@ struct zi_dev {
     uint16_t *in_bf16;
 };
 
-/* Convert the mmapped F32 weights in bounded chunks (2 MB of host scratch),
- * rounding to nearest even like the device conversion. Concatenation is on
- * the device, so neither QKV nor the gate/up weights need a host-sized copy. */
+/* dst[i] = bf16(src[i]), rounding to nearest even like the device conversion. */
+static void zi_bf16(uint16_t *dst, const float *src, size_t n) {
+    for (size_t i = 0; i < n; i++) {
+        uint32_t bits;
+        memcpy(&bits, src + i, sizeof(bits));
+        bits += 0x7fff + ((bits >> 16) & 1);
+        dst[i] = (uint16_t)(bits >> 16);
+    }
+}
+
+/* Convert the mmapped F32 weights in bounded chunks (2 MB of host scratch).
+ * Concatenation is on the device, so neither QKV nor the gate/up weights
+ * need a host-sized copy. */
 static uint16_t *zi_dev_rows(const float *const *parts, const size_t *counts, int n) {
     size_t total = 0;
     for (int i = 0; i < n; i++) {
@@ -1733,12 +1749,7 @@ static uint16_t *zi_dev_rows(const float *const *parts, const size_t *counts, in
         for (size_t j = 0; j < counts[i]; j += chunk) {
             size_t len = counts[i] - j;
             if (len > chunk) len = chunk;
-            for (size_t k = 0; k < len; k++) {
-                uint32_t bits;
-                memcpy(&bits, parts[i] + j + k, sizeof(bits));
-                bits += 0x7fff + ((bits >> 16) & 1);
-                tmp[k] = (uint16_t)(bits >> 16);
-            }
+            zi_bf16(tmp, parts[i] + j, len);
             iris_dev_upload(out + off + j, tmp, len * sizeof(uint16_t));
         }
         off += counts[i];
@@ -1758,16 +1769,45 @@ static float *zi_dev_f32(const float *p, size_t n) {
     return out;
 }
 
-static int zi_dev_load_block(zi_transformer_t *tf, zi_dev_block_t *d, const zi_block_t *b) {
+/* Elements of the qkv, proj, ff_in and ff_out matrices of a block. */
+static size_t zi_dev_layer_elems(const zi_transformer_t *tf) {
+    size_t dd = (size_t)tf->dim * tf->dim, md = (size_t)tf->dim * tf->ffn_dim;
+    return 4 * dd + 3 * md;
+}
+
+/* Move the four big matrices of a block to host memory, in the slot layout. */
+static int zi_dev_stream(zi_transformer_t *tf, zi_dev_block_t *d, const zi_block_t *b) {
+    size_t dd = (size_t)tf->dim * tf->dim, md = (size_t)tf->dim * tf->ffn_dim;
+    const float *w[] = {b->attn_q_weight, b->attn_k_weight, b->attn_v_weight,
+                        b->attn_out_weight, b->ffn_w1, b->ffn_w3, b->ffn_w2};
+    size_t n[] = {dd, dd, dd, dd, md, md, md};
+    d->host = malloc(zi_dev_layer_elems(tf) * sizeof(uint16_t));
+    if (!d->host) return 0;
+    for (size_t i = 0, off = 0; i < 7; off += n[i++]) zi_bf16(d->host + off, w[i], n[i]);
+    iris_dev_free(d->qkv);
+    iris_dev_free(d->proj);
+    iris_dev_free(d->ff_in);
+    iris_dev_free(d->ff_out);
+    d->qkv = d->proj = d->ff_in = d->ff_out = NULL;
+    return 1;
+}
+
+static int zi_dev_load_block(zi_transformer_t *tf, zi_dev_block_t *d, const zi_block_t *b,
+                             int stream) {
     size_t dim = tf->dim, dd = dim * dim, md = dim * tf->ffn_dim;
     const float *qkv[] = {b->attn_q_weight, b->attn_k_weight, b->attn_v_weight};
     size_t qkv_n[] = {dd, dd, dd};
     const float *ff[] = {b->ffn_w1, b->ffn_w3};
     size_t ff_n[] = {md, md};
-    d->qkv = zi_dev_rows(qkv, qkv_n, 3);
-    d->proj = zi_dev_bf16(b->attn_out_weight, dd);
-    d->ff_in = zi_dev_rows(ff, ff_n, 2);
-    d->ff_out = zi_dev_bf16(b->ffn_w2, md);
+    if (stream) {
+        if (!zi_dev_stream(tf, d, b)) return 0;
+    } else {
+        d->qkv = zi_dev_rows(qkv, qkv_n, 3);
+        d->proj = zi_dev_bf16(b->attn_out_weight, dd);
+        d->ff_in = zi_dev_rows(ff, ff_n, 2);
+        d->ff_out = zi_dev_bf16(b->ffn_w2, md);
+        if (!d->qkv || !d->proj || !d->ff_in || !d->ff_out) return 0;
+    }
     d->norm_q = zi_dev_f32(b->attn_norm_q, tf->head_dim);
     d->norm_k = zi_dev_f32(b->attn_norm_k, tf->head_dim);
     d->attn_norm1 = zi_dev_f32(b->attn_norm1, dim);
@@ -1779,29 +1819,42 @@ static int zi_dev_load_block(zi_transformer_t *tf, zi_dev_block_t *d, const zi_b
         d->mod_bias = zi_dev_f32(b->adaln_bias, 4 * dim);
         if (!d->mod || !d->mod_bias) return 0;
     }
-    return d->qkv && d->proj && d->ff_in && d->ff_out && d->norm_q && d->norm_k &&
-           d->attn_norm1 && d->attn_norm2 && d->ffn_norm1 && d->ffn_norm2;
+    return d->norm_q && d->norm_k && d->attn_norm1 && d->attn_norm2 && d->ffn_norm1 && d->ffn_norm2;
+}
+
+static void zi_dev_free_block(zi_dev_block_t *b) {
+    iris_dev_free(b->qkv);
+    iris_dev_free(b->proj);
+    iris_dev_free(b->ff_in);
+    iris_dev_free(b->ff_out);
+    iris_dev_free(b->mod);
+    iris_dev_free(b->mod_bias);
+    iris_dev_free(b->norm_q);
+    iris_dev_free(b->norm_k);
+    iris_dev_free(b->attn_norm1);
+    iris_dev_free(b->attn_norm2);
+    iris_dev_free(b->ffn_norm1);
+    iris_dev_free(b->ffn_norm2);
+    free(b->host);
+    memset(b, 0, sizeof(*b));
 }
 
 static void zi_dev_free_blocks(zi_dev_block_t *b, int n) {
     if (!b) return;
-    for (int i = 0; i < n; i++) {
-        iris_dev_free(b[i].qkv);
-        iris_dev_free(b[i].proj);
-        iris_dev_free(b[i].ff_in);
-        iris_dev_free(b[i].ff_out);
-        iris_dev_free(b[i].mod);
-        iris_dev_free(b[i].mod_bias);
-        iris_dev_free(b[i].norm_q);
-        iris_dev_free(b[i].norm_k);
-        iris_dev_free(b[i].attn_norm1);
-        iris_dev_free(b[i].attn_norm2);
-        iris_dev_free(b[i].ffn_norm1);
-        iris_dev_free(b[i].ffn_norm2);
-    }
+    for (int i = 0; i < n; i++) zi_dev_free_block(&b[i]);
     free(b);
 }
 
+/* Stream the last resident main layer, freeing its GPU memory.
+ * Returns 0 if no layer is left on the GPU. */
+static int zi_dev_evict(zi_transformer_t *tf) {
+    struct zi_dev *d = tf->dev;
+    int i = d->resident - 1;
+    if (i < 0 || !zi_dev_stream(tf, &d->layers[i], &tf->layers[i])) return 0;
+    d->resident = i;
+    return 1;
+}
+
 static void zi_dev_free_activations(struct zi_dev *d) {
     iris_dev_free(d->x); d->x = NULL;
     iris_dev_free(d->out); d->out = NULL;
@@ -1820,6 +1873,7 @@ static void zi_dev_free(zi_transformer_t *tf) {
     zi_dev_free_blocks(d->noise, tf->n_refiner);
     zi_dev_free_blocks(d->context, tf->n_refiner);
     zi_dev_free_blocks(d->layers, tf->n_layers);
+    iris_dev_free(d->slot);
     iris_dev_free(d->t_lin1);
     iris_dev_free(d->t_lin2);
     iris_dev_free(d->t_bias1);
@@ -1873,16 +1927,28 @@ static int zi_dev_load(zi_transformer_t *tf) {
     d->context = calloc(tf->n_refiner, sizeof(*d->context));
     d->layers = calloc(tf->n_layers, sizeof(*d->layers));
     if ((tf->n_refiner && (!d->noise || !d->context)) || (tf->n_layers && !d->layers)) goto fail;
-    for (int i = 0; i < tf->n_refiner; i++) {
-        if (!zi_dev_load_block(tf, &d->noise[i], &tf->noise_refiner[i]) ||
-            !zi_dev_load_block(tf, &d->context[i], &tf->context_refiner[i])) goto fail;
-    }
-    for (int i = 0; i < tf->n_layers; i++)
-        if (!zi_dev_load_block(tf, &d->layers[i], &tf->layers[i])) goto fail;
     d->vec = iris_dev_alloc((256 + mid + ad) * sizeof(float));
     d->vec_bf16 = iris_dev_alloc((256 + mid + ad) * sizeof(uint16_t));
     d->mods = iris_dev_alloc(((size_t)(tf->n_refiner + tf->n_layers) * 4 + 1) * dim * sizeof(float));
     if (!d->vec || !d->vec_bf16 || !d->mods) goto fail;
+    for (int i = 0; i < tf->n_refiner; i++) {
+        if (!zi_dev_load_block(tf, &d->noise[i], &tf->noise_refiner[i], 0) ||
+            !zi_dev_load_block(tf, &d->context[i], &tf->context_refiner[i], 0)) goto fail;
+    }
+    /* When the GPU is full, this layer and the next ones are streamed. If
+     * even their small vectors do not fit, resident layers make room. */
+    int stream = 0;
+    for (int i = 0; i < tf->n_layers; ) {
+        if (zi_dev_load_block(tf, &d->layers[i], &tf->layers[i], stream)) {
+            d->resident += !stream;
+            i++;
+            continue;
+        }
+        zi_dev_free_block(&d->layers[i]);
+        if (stream && !zi_dev_evict(tf)) goto fail;
+        stream = 1;
+    }
+    d->shown = tf->n_layers;
     iris_dev_sync();
     return 1;
 fail:
@@ -1916,6 +1982,10 @@ static int zi_dev_ensure(zi_transformer_t *tf, int seq, size_t in_elems) {
         if (!d->in_f32 || !d->in_bf16) return 0;
         d->in_elems = in_elems;
     }
+    if (d->resident < tf->n_layers && !d->slot) {
+        d->slot = iris_dev_alloc(zi_dev_layer_elems(tf) * sizeof(uint16_t));
+        if (!d->slot) return 0;
+    }
     return 1;
 }
 
@@ -1968,7 +2038,14 @@ static float *zi_forward_dev(zi_transformer_t *tf, const float *latent, int late
     int seq = ip + cp, half = tf->head_dim / 2;
     size_t in_elems = (size_t)img_seq * patch;
     if (in_elems < (size_t)cap_seq * tf->cap_feat_dim) in_elems = (size_t)cap_seq * tf->cap_feat_dim;
-    if (!zi_dev_ensure(tf, seq, in_elems)) goto fail;
+    /* Make room for the activations at this size, streaming layers if needed. */
+    while (!zi_dev_ensure(tf, seq, in_elems))
+        if (!zi_dev_evict(tf)) goto fail;
+    if (d->resident != d->shown) {
+        fprintf(stderr, "Z-Image: %d of %d layers on the GPU, the others uploaded at each step\n",
+                d->resident, tf->n_layers);
+        d->shown = d->resident;
+    }
     float *tokens = malloc((size_t)img_seq * patch * sizeof(float));
     float *output = calloc((size_t)tf->in_channels * latent_h * latent_w, sizeof(float));
     int *pos = calloc((size_t)seq * 3, sizeof(int));
@@ -2054,8 +2131,17 @@ static float *zi_forward_dev(zi_transformer_t *tf, const float *latent, int late
     iris_dev_sync();
     double context_time = zi_time_ms() - start;
     start = zi_time_ms();
+    size_t dd = (size_t)dim * dim, md = (size_t)dim * tf->ffn_dim;
     for (int i = 0; i < tf->n_layers; i++) {
-        zi_dev_block(tf, &d->layers[i], d->x, seq, 0, d->mods + (size_t)(tf->n_refiner + i) * 4 * dim,
+        zi_dev_block_t b = d->layers[i];
+        if (b.host) {
+            iris_dev_upload(d->slot, b.host, zi_dev_layer_elems(tf) * sizeof(uint16_t));
+            b.qkv = d->slot;
+            b.proj = b.qkv + 3 * dd;
+            b.ff_in = b.proj + dd;
+            b.ff_out = b.ff_in + 2 * md;
+        }
+        zi_dev_block(tf, &b, d->x, seq, 0, d->mods + (size_t)(tf->n_refiner + i) * 4 * dim,
                      ip == img_seq && cp == cap_seq ? NULL : d->mask);
         if (iris_substep_callback) iris_substep_callback(IRIS_SUBSTEP_SINGLE_BLOCK, i, tf->n_layers);
     }

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