Removing Large Zero Tensors from Ideogram 4 Conditioning
Cutting peak GPU memory by 9.46 GiB at 2048², with pixel-identical images.
Baseline

Optimized

Baseline and optimized output. The panels intentionally use the same image bytes because the measured output is pixel-identical.
After exploring masked diffusion models, I started looking at image models and profiled Ideogram 4 NF4, an open-source model. What I found was that at 1024², the pipeline built a (1, 4608, 53248) positive-conditioning tensor, but only 512 of those rows contained text features; the other 4,096 were zero. Classifier-free guidance (CFG) also introduced a separate (1, 4096, 53248) negative-conditioning tensor that was entirely zero.
The first runtime optimization eliminates redundant computation in text-conditioning reducing peak GPU memory while preserving runtime and pixel-idential output.
Looking further in the model, it was observed that the prompt stayed fixed. At every denoising step, the projection of the text-conditioning would be recomputated. The second optimization involves caching the NF4 projection of the text-conditioning prompt, removing this unecessary computation.
This post assumes basic familiarity with diffusion transformers.
Conditioning a Denoising Step
Ideogram updates a noisy image latent over several denoising steps. A Qwen text encoder converts the fixed prompt into one 53,248-dimensional feature vector per text position. For text positions, its output is
The image transformer has a hidden width of 4,608, so a learned linear layer projects each Qwen row into that space:
At 1024² resolution, the transformer also processes 4,096 image positions. Combined with the 512-token prompt, the packed sequence has 4,608 positions:
The sequence length and hidden width happen to be the same number here. They are different dimensions.
Most of the Conditioning Tensor Was Zero
The Qwen encoder produces features only for text positions. The Diffusers path still expanded the conditioning tensor to match the full text-and-image sequence:
where contains one zero row for each image position. At 1024², 4,096 of the 4,608 positive-branch rows were empty. Classifier-free guidance added a negative branch of another 4,096 zero rows.
The projection has a bias, so a zero row does not stay zero:
The original code masked the image positions after projection. It therefore built zero rows, projected them into identical bias vectors, and then discarded those outputs. Since the linear layer handles each row independently, the empty image rows cannot affect the projection of the real text rows.
The wide tensors are built in fp32 before the cast to bf16. Together, the two inputs occupy
The allocator trace matched that calculation: active memory rose from 14.998 GiB to 16.725 GiB during construction. The positive and negative projections then took a median 9.969 ms and 8.374 ms. A 20-step generation repeated both on the same prompt twenty times.
I changed the path to keep only the real text rows, so the wide tensor is (1, 512, 53248). The fully masked negative conditioning became None instead of a materialized tensor. Nothing else moved: sampler, schedule, timesteps, latents, weights, guidance scale, attention implementation, and branch order all stayed as they were.
Optimization 2: Moving text-conditioning projection
The second thing I noticed was the text features and projection do not change during generation. As a result, the prompt only needs to be projected once before denoising begins:
For the 512-token run, each term maps to a concrete part of the pipeline:
| Term | Meaning in the pipeline |
|---|---|
Qwen text features for the prompt, shape (1, 512, 53248) | |
llm_cond_norm, which normalizes the text features before projection | |
The learned NF4 llm_cond_proj, which maps each row from width 53,248 to 4,608 | |
The projected text conditioning consumed by the diffusion transformer, shape (1, 512, 4608) |
The original 53,248-wide text features is an input to the normalization and projection function, which then becomes the input to the diffusion transformer. Since the prompt and projection stay fixed during denoising, stays fixed too. By computing it once and caching the roughly 4.5 MiB bf16 result, I can write it into the packed hidden states at each step.
Shape-Dependent NF4 Numerics
The obvious next step was to normalize and project only the real text rows. Mathematically, these outputs should match the first rows of the original -row projection. In practice, they were not bitwise identical.
I first normalized and projected only the real rows. Mathematically, those outputs should match the first rows of the original -row call. They did not match bitwise.
The projection is a BitsAndBytes Linear4bit layer with NF4 weights. Its matrix multiplication has the form
where , , and is the input row count. The baseline used . The smaller call used .
BitsAndBytes chooses different kernels for different matrix shapes. Projecting only the 512 real rows changed the order of dequantization and accumulation, even though the math was nominally the same.
The direct 512-row projection differed from the 4,608-row reference by up to 0.015625. After one denoising step, the guided velocity differed by as much as 0.546875. The outputs were still extremely close, with 0.9997085 cosine similarity, but no longer pixel identical.
I tested several smaller padding sizes, but none preserved exactness consistently across prompt lengths. Padding to 2,048 rows, for example, matched the reference for 12 of 18 full generations and diverged on the remaining six. The only configuration that reproduced the reference reliably across all prompts was the original row count:
Preserving the Reference Shape
The working path builds a temporary scratch tensor with exactly rows. Its first rows contain the normalized text features; the remaining rows repeat them until the scratch reaches the original packed length. The filler values cannot affect the real outputs because the projection handles rows independently. They only make BitsAndBytes see the baseline GEMM shape.
The layer runs once at that shape, in a median 14.762 ms. I keep the first projected rows, free the scratch and unused outputs, and reuse the (B, T, 4608) result throughout denoising.
real text rows
-> repeat to the original packed row count
-> run the NF4 projection once
-> keep the first T outputs
-> reuse them during denoisingThe wide zero-padded conditioning disappears from the per-step path. The one-time projection keeps the reference arithmetic.
Results
I benchmarked batch-size-one inference on Modal with the prompt, seed, schedule, dtype, and branch order held constant. Reported latencies are the median of three runs after one warm-up run. The environment used Diffusers 7685bffe, model revision 1874bc70, PyTorch 2.13.0, and BitsAndBytes 0.50.0, with all model weights kept on GPU. I did not test CPU offloading or batch sizes greater than one.
The memory savings increase with image resolution because the removed Qwen-width rows scale with the number of image positions.
| Resolution | Image positions | Baseline peak | Optimized peak | Saved | Step change |
|---|---|---|---|---|---|
| 512² | 1,024 | 16.041003 GiB | 15.722261 GiB | 0.318742 GiB | −2.04% |
| 1024² | 4,096 | 18.189322 GiB | 16.053430 GiB | 2.135892 GiB | −2.27% |
| 2048² | 16,384 | 26.835333 GiB | 17.378107 GiB | 9.457227 GiB | −1.26% |
The measured drops are larger than the retained bf16 zero rows alone predict because the fp32 construction and per-step projection temporaries also disappear. Runtime improves by only 1 to 3 percent because the transformer still dominates each step.
Across 18 full 1024² generations, covering three prompt lengths, three seeds, and both Turbo-12 and Default-20, the final latents and images were bitwise identical. Tesseract also returned identical text and word boxes, which matters for a model built to render text. This exactness result is scoped to the pinned BitsAndBytes stack and the tested shapes.
Full-generation peak memory drops by 0.792383 GiB on the short prompt and 1.620291 GiB on the near-maximum one. The whole generation runs 2 to 3% faster.
References
- QLoRA: Efficient Finetuning of Quantized LLMs, which introduced NF4 (arXiv:2305.14314)
- ideogram-oss/ideogram4
- huggingface/diffusers
Measurements were taken on Modal L40S, A100-80GB, and A10 instances