Ideogram4: allow the text conditioning projection to be hoisted out of the denoising loop

[lsnchow]

I was profiling Ideogram 4 inference on a 24 GB card and found that most of the text conditioning
work repeats on every denoising step.

Ideogram4Transformer2DModel.forward normalizes and projects encoder_hidden_states each step:

encoder_hidden_states = self.llm_cond_norm(encoder_hidden_states)
encoder_hidden_states = self.llm_cond_proj(encoder_hidden_states) * llm_token_mask

But encoder_hidden_states is fixed for the whole generation, and the projection runs at the full
packed width before masking most of the result away. The tensor is 53,248 rows wide and 512 carry
text, so at 2048² that is 16,384 of 16,896 rows projected and discarded. Classifier-free guidance
adds a second branch that is zero throughout. Under a 4 bit llm_cond_proj that is a wide NF4 GEMM
per step over rows that cannot affect the output.

Hoisting it out of the loop, with bitsandbytes NF4 against 7685bffe8904:

	baseline	hoisted
peak allocation, 2048² (A100)	26.835333 GiB	17.378107 GiB	−9.457 GiB
peak allocation, 1024² (L40S)	18.189322 GiB	16.053430 GiB	−2.136 GiB

Byte identical output, same SHA-256 on latents, pixels, and OCR, across 18 generations over three
prompt lengths and three seeds. At 2048² the unmodified path OOMs on an A10's 24 GiB. The hoisted
path completes.

Why It Needs an API Change

I tried to avoid one. The caller builds the (1, 16896, 53248) tensor before forward is entered,
so the transformer cannot unspend it. And skipping the masked rows internally changes the GEMM's M,
which changes which NF4 kernel runs: projecting only the 512 real rows changes pixels. I measured
that. Preserving the row count is what keeps it exact.

Proposed Shape

One public method and one keyword only flag, both opt-in, default path untouched:

projected = transformer.project_text_conditioning(
    text_features[:, indicator[0] == LLM_TOKEN_INDICATOR],
    # Pads back to the row count the default path would have used, so a quantized llm_cond_proj
    # selects the same kernel and the result stays bitwise identical.
    minimum_projection_rows=text_features.shape[0] * text_features.shape[1],
)
# then per step: transformer(..., encoder_hidden_states=projected, encoder_hidden_states_projected=True)

An unconditional branch with no text positions can pass encoder_hidden_states=None instead of a
wide all zero tensor.

Would you take a change of this shape? Implemented against 7685bffe8904 with tests and docs,
make style / make quality / make fix-copies clean, and the model's tests passing on CPU:
https://github.com/lsnchow/diffusers/tree/ideogram4-project-text-conditioning-once

Two things I would rather hear your view on. The opt-in path is not torch.compile(fullgraph=True)
compatible, because rejecting a ragged batch means reading indicator. The default path still
compiles. And minimum_projection_rows is an odd argument for a public signature, but without it
the change is no longer exact, which removes the reason to make it.

The code was written by an AI agent, per the agentic contribution guidelines. I have run
self-review and can post the notes.

@yiyixuxu @DN6

[Mahnoor-Zaffar]

Hey! Thanks for reporting this.

The most likely root cause of this issue is the repeated normalization and projection of encoder_hidden_states within the forward method. Specifically, the lines:

encoder_hidden_states = self.llm_cond_norm(encoder_hidden_states)
encoder_hidden_states = self.llm_cond_proj(encoder_hidden_states) * llm_token_mask

are executed for every denoising step, despite the input being constant. This can be found in the Ideogram4Transformer2DModel class, particularly in the forward function.

To resolve this, the text conditioning projection should be hoisted out of the denoising loop, allowing it to be computed once before the loop begins. This change should maintain the integrity of the output while reducing memory usage significantly.

If no one is currently working on this, I'd be happy to take a look and put together a fix.

[lsnchow]

I’ve actually already started working on an implementation. Happy to open a PR soon if the maintainers are interested.