black-forest-labs

/

FLUX.2-klein-9B

like 1.22k

Follow

Black Forest Labs 40.4k

Issue

FLUX.2 Klein pads every prompt to max_sequence_length=512:

inputs = tokenizer(
    text,
    padding="max_length",
    truncation=True,
    max_length=max_sequence_length,
    return_tensors="pt",
)

The resulting attention_mask is correctly passed to Qwen3ForCausalLM, so padding tokens are masked inside Qwen3 self-attention.

However, FLUX.2 then concatenates intermediate Qwen3 hidden states:

out = torch.stack(
    [output.hidden_states[k] for k in hidden_states_layers],
    dim=1,
)

prompt_embeds = out.permute(0, 2, 1, 3).reshape(
    batch_size, seq_len, num_channels * hidden_dim
)

The output still has length 512, including hidden states corresponding to padded positions.

After this step:

• the padding mask is not returned;
• padded hidden states are not removed or zeroed;
• position IDs are created for all 512 positions;
• no text padding mask is passed to the FLUX.2 transformer.

Therefore, for a short prompt, FLUX.2 joint attention appears to receive something like:

20 valid text embeddings
492 padding-derived text embeddings

All 512 positions then participate in text-image joint attention.

Question

Is this behavior intentional?

Specifically:

1. Were FLUX.2 checkpoints trained with the padded Qwen3 hidden states as learned null/background conditioning tokens?
2. Or should the valid Qwen3 attention mask also be applied inside FLUX.2 joint attention?
3. Would trimming prompt_embeds to the actual prompt length be incorrect for the released checkpoint?

If the fixed-length unmasked sequence is intentional, it would be helpful to document this behavior.