Removing Large Zero Tensors from Ideogram 4 Conditioning

Cutting peak GPU memory by 9.46 GiB at 2048², with pixel-identical images.

Lucas ChowJuly 31, 2026 · 7 min · source

Baseline

A golden retriever wearing sunglasses on the hood of a red convertible beside a coastal road at sunset

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 TT text positions, its output is

XRB×T×53248.X \in \mathbb{R}^{B \times T \times 53248}.

The image transformer has a hidden width of 4,608, so a learned linear layer projects each Qwen row into that space:

P(X)=XW+b,WR4608×53248.P(X) = XW^\top + b, \qquad W \in \mathbb{R}^{4608 \times 53248}.

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:

C+=[X0I],C_+ = \begin{bmatrix} X \\ 0_I \end{bmatrix},

where 0I0_I 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:

P(0)=b.P(0) = b.

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

(4608+4096)×53248×41.7266 GiB.(4608 + 4096) \times 53248 \times 4 \approx 1.7266\ \text{GiB}.

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:

Z=P(N(X)).Z = P(N(X)).

For the 512-token run, each term maps to a concrete part of the pipeline:

TermMeaning in the pipeline
XXQwen text features for the prompt, shape (1, 512, 53248)
N()N(\cdot)llm_cond_norm, which normalizes the text features before projection
P()P(\cdot)The learned NF4 llm_cond_proj, which maps each row from width 53,248 to 4,608
ZZThe 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, ZZ 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 (T)(T) real text rows. Mathematically, these outputs should match the first (T)(T) rows of the original (T+I)(T + I)-row projection. In practice, they were not bitwise identical.

I first normalized and projected only the TT real rows. Mathematically, those outputs should match the first TT rows of the original (T+I)(T + I)-row call. They did not match bitwise.

The projection is a BitsAndBytes Linear4bit layer with NF4 weights. Its matrix multiplication has the form

AM×KWK×N,A_{M \times K} W_{K \times N},

where K=53248K = 53248, N=4608N = 4608, and MM is the input row count. The baseline used M=T+IM = T + I. The smaller call used M=TM = T.

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:

M=T+I.M = T + I.

Preserving the Reference Shape

The working path builds a temporary scratch tensor with exactly T+IT + I rows. Its first TT 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 TT 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 denoising

The 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.

ResolutionImage positionsBaseline peakOptimized peakSavedStep change
512²1,02416.041003 GiB15.722261 GiB0.318742 GiB−2.04%
1024²4,09618.189322 GiB16.053430 GiB2.135892 GiB−2.27%
2048²16,38426.835333 GiB17.378107 GiB9.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

Measurements were taken on Modal L40S, A100-80GB, and A10 instances