[Qwen-image] encoder_hidden_states_mask is not used

@sayakpaul

Describe the bug

During text encoding stage, the text encoder will produce encoder_hidden_states_mask to mask some regions. It usually happens during training with batch size more than 1, to make sure that the produce text embedding has the same sequence length.

However, when I check the following code, it seems weird that encoder_hidden_states_mask is not used even though it has been passed to the function.

It indeed works when batch=1. However, I feel it is indeed a bug if we train Qwen-image with batch size more than 1.

Reproduction

class QwenDoubleStreamAttnProcessor2_0:
    """
    Attention processor for Qwen double-stream architecture, matching DoubleStreamLayerMegatron logic. This processor
    implements joint attention computation where text and image streams are processed together.
    """

    _attention_backend = None

    def __init__(self):
        if not hasattr(F, "scaled_dot_product_attention"):
            raise ImportError(
                "QwenDoubleStreamAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
            )

    def __call__(
        self,
        attn: Attention,
        hidden_states: torch.FloatTensor,  # Image stream
        encoder_hidden_states: torch.FloatTensor = None,  # Text stream
        encoder_hidden_states_mask: torch.FloatTensor = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        image_rotary_emb: Optional[torch.Tensor] = None,
    ) -> torch.FloatTensor:
        if encoder_hidden_states is None:
            raise ValueError("QwenDoubleStreamAttnProcessor2_0 requires encoder_hidden_states (text stream)")

        seq_txt = encoder_hidden_states.shape[1]

        # Compute QKV for image stream (sample projections)
        img_query = attn.to_q(hidden_states)
        img_key = attn.to_k(hidden_states)
        img_value = attn.to_v(hidden_states)

        # Compute QKV for text stream (context projections)
        txt_query = attn.add_q_proj(encoder_hidden_states)
        txt_key = attn.add_k_proj(encoder_hidden_states)
        txt_value = attn.add_v_proj(encoder_hidden_states)

        # Reshape for multi-head attention
        img_query = img_query.unflatten(-1, (attn.heads, -1))
        img_key = img_key.unflatten(-1, (attn.heads, -1))
        img_value = img_value.unflatten(-1, (attn.heads, -1))

        txt_query = txt_query.unflatten(-1, (attn.heads, -1))
        txt_key = txt_key.unflatten(-1, (attn.heads, -1))
        txt_value = txt_value.unflatten(-1, (attn.heads, -1))

        # Apply QK normalization
        if attn.norm_q is not None:
            img_query = attn.norm_q(img_query)
        if attn.norm_k is not None:
            img_key = attn.norm_k(img_key)
        if attn.norm_added_q is not None:
            txt_query = attn.norm_added_q(txt_query)
        if attn.norm_added_k is not None:
            txt_key = attn.norm_added_k(txt_key)

        # Apply RoPE
        if image_rotary_emb is not None:
            img_freqs, txt_freqs = image_rotary_emb
            img_query = apply_rotary_emb_qwen(img_query, img_freqs, use_real=False)
            img_key = apply_rotary_emb_qwen(img_key, img_freqs, use_real=False)
            txt_query = apply_rotary_emb_qwen(txt_query, txt_freqs, use_real=False)
            txt_key = apply_rotary_emb_qwen(txt_key, txt_freqs, use_real=False)

        # Concatenate for joint attention
        # Order: [text, image]
        joint_query = torch.cat([txt_query, img_query], dim=1)
        joint_key = torch.cat([txt_key, img_key], dim=1)
        joint_value = torch.cat([txt_value, img_value], dim=1)

        # Compute joint attention
        joint_hidden_states = dispatch_attention_fn(
            joint_query,
            joint_key,
            joint_value,
            attn_mask=attention_mask,
            dropout_p=0.0,
            is_causal=False,
            backend=self._attention_backend,
        )

        # Reshape back
        joint_hidden_states = joint_hidden_states.flatten(2, 3)
        joint_hidden_states = joint_hidden_states.to(joint_query.dtype)

        # Split attention outputs back
        txt_attn_output = joint_hidden_states[:, :seq_txt, :]  # Text part
        img_attn_output = joint_hidden_states[:, seq_txt:, :]  # Image part

        # Apply output projections
        img_attn_output = attn.to_out[0](img_attn_output)
        if len(attn.to_out) > 1:
            img_attn_output = attn.to_out[1](img_attn_output)  # dropout

        txt_attn_output = attn.to_add_out(txt_attn_output)

        return img_attn_output, txt_attn_output

Logs

System Info

A100

Who can help?

@sayakpaul

@naykun

Cc: @naykun

Thanks for pointing this out. I would have expected that the encoder_hidden_states mask is translated to an attention mask ask as it is done in other pipeline (like Chroma). So maybe it is used, but then the encoder_hidden_states_mask parameter here should be removed

I've checked, it's actually not used. In Chroma, a similar 1D mask is translated here

diffusers/src/diffusers/models/transformers/transformer_chroma.py

Line 324 in fc337d5

    
           attention_mask = attention_mask[:, None, None, :] * attention_mask[:, None, :, None]

into a 2D attention mask as the attention processor expects.

In Qwen, no attention mask is passed...

diffusers/src/diffusers/models/transformers/transformer_qwenimage.py

Line 436 in fc337d5

...and the hidden_states_mask is unused as OP has shown.

Bug workaround:

            #bug workaround for https://github.com/huggingface/diffusers/issues/12294
            image_attention_mask=torch.ones((batch_size, image_seq_len), dtype=torch.bool, device=latent_image.device)
            attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1)
            attention_mask_2d = attention_mask[:, None, None, :] * attention_mask[:, None, :, None]

                noise_pred = transformer(

...

                    attention_kwargs = {
                        "attention_mask": attention_mask_2d,
                    },

Something similar is probably also the fix to this problem

@naykun

Cc: @naykun

As a general recommendation, we always follow the original implementation code and many a times, it's only provided for the context of just inference. So, if you want to fix, we'd appreciate if you also opened PRs and tagged the original authors of a particular model.

Cc: @naykun

As a general recommendation, we always follow the original implementation code and many a times, it's only provided for the context of just inference. So, if you want to fix, we'd appreciate if you also opened PRs and tagged the original authors of a particular model.

Note that this can also affect inference: It's common practice for a CFG model, to infer the conditional and the unconditional in the same batch. The unconditional encoding is usually much shorter and if not correctly masked, contains visible padding token encodings. I've tried that and it makes a real difference. Without the workaround above, the image is blurry and/or has less detail.

The diffusers pipeline itself is not affected, because the transformer is called twice separately with batch size 1.

I have a worked-out solution Here's the diff:

Unfold

diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py
index 977f2790a..8758f608b 100644
--- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py
+++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py
@@ -546,6 +546,7 @@ class QwenImageEditPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
         callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
         callback_on_step_end_tensor_inputs: List[str] = ["latents"],
         max_sequence_length: int = 512,
+        batch_cfg: bool = False,
     ):
         r"""
         Function invoked when calling the pipeline for generation.
@@ -712,6 +713,14 @@ class QwenImageEditPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
                 num_images_per_prompt=num_images_per_prompt,
                 max_sequence_length=max_sequence_length,
             )
+            if batch_cfg:
+                target_len = max(negative_prompt_embeds.size(1), prompt_embeds.size(1))
+                negative_prompt_embeds      = self._pad_to_len(negative_prompt_embeds, target_len, pad_value=0.0)
+                prompt_embeds               = self._pad_to_len(prompt_embeds, target_len, pad_value=0.0)
+                negative_prompt_embeds_mask = self._pad_to_len(negative_prompt_embeds_mask, target_len, pad_value=0)
+                prompt_embeds_mask          = self._pad_to_len(prompt_embeds_mask, target_len, pad_value=0)
+                prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+                prompt_embeds_mask = torch.cat([negative_prompt_embeds_mask, prompt_embeds_mask], dim=0)
 
         # 4. Prepare latent variables
         num_channels_latents = self.transformer.config.in_channels // 4
@@ -732,7 +741,7 @@ class QwenImageEditPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
                 (1, calculated_height // self.vae_scale_factor // 2, calculated_width // self.vae_scale_factor // 2),
             ]
         ] * batch_size
-
+    
         # 5. Prepare timesteps
         sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
         image_seq_len = latents.shape[1]
@@ -771,9 +780,10 @@ class QwenImageEditPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
             self._attention_kwargs = {}
 
         txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
-        negative_txt_seq_lens = (
-            negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None
-        )
+        if not batch_cfg:
+            negative_txt_seq_lens = (
+                negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None
+            )
 
         # 6. Denoising loop
         self.scheduler.set_begin_index(0)
@@ -787,9 +797,14 @@ class QwenImageEditPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
                 latent_model_input = latents
                 if image_latents is not None:
                     latent_model_input = torch.cat([latents, image_latents], dim=1)
+                if batch_cfg:
+                    latent_model_input = torch.cat([latent_model_input] * 2)
 
                 # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-                timestep = t.expand(latents.shape[0]).to(latents.dtype)
+                if batch_cfg:
+                    timestep = t.expand(latent_model_input.shape[0]).to(latent_model_input.dtype)
+                else:
+                    timestep = t.expand(latents.shape[0]).to(latents.dtype)
                 with self.transformer.cache_context("cond"):
                     noise_pred = self.transformer(
                         hidden_states=latent_model_input,
@@ -805,19 +820,22 @@ class QwenImageEditPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
                     noise_pred = noise_pred[:, : latents.size(1)]
 
                 if do_true_cfg:
-                    with self.transformer.cache_context("uncond"):
-                        neg_noise_pred = self.transformer(
-                            hidden_states=latent_model_input,
-                            timestep=timestep / 1000,
-                            guidance=guidance,
-                            encoder_hidden_states_mask=negative_prompt_embeds_mask,
-                            encoder_hidden_states=negative_prompt_embeds,
-                            img_shapes=img_shapes,
-                            txt_seq_lens=negative_txt_seq_lens,
-                            attention_kwargs=self.attention_kwargs,
-                            return_dict=False,
-                        )[0]
-                    neg_noise_pred = neg_noise_pred[:, : latents.size(1)]
+                    if not batch_cfg:
+                        with self.transformer.cache_context("uncond"):
+                            neg_noise_pred = self.transformer(
+                                hidden_states=latent_model_input,
+                                timestep=timestep / 1000,
+                                guidance=guidance,
+                                encoder_hidden_states_mask=negative_prompt_embeds_mask,
+                                encoder_hidden_states=negative_prompt_embeds,
+                                img_shapes=img_shapes,
+                                txt_seq_lens=negative_txt_seq_lens,
+                                attention_kwargs=self.attention_kwargs,
+                                return_dict=False,
+                            )[0]
+                        neg_noise_pred = neg_noise_pred[:, : latents.size(1)]
+                    else:
+                        neg_noise_pred, noise_pred = noise_pred.chunk(2)
                     comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
 
                     cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
@@ -874,3 +892,23 @@ class QwenImageEditPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
             return (image,)
 
         return QwenImagePipelineOutput(images=image)
+
+    @staticmethod
+    def _pad_to_len(x, target_len, pad_value=0.0):
+        # x: [B, S, D] or [B, S]
+        if x.dim() == 3:   # embeds
+            B, S, D = x.shape
+            if S == target_len:
+                return x
+            out = x.new_full((B, target_len, D), pad_value)
+            out[:, :S, :] = x
+            return out
+        elif x.dim() == 2: # mask
+            B, S = x.shape
+            if S == target_len:
+                return x
+            out = x.new_zeros((B, target_len), dtype=x.dtype)
+            out[:, :S] = x
+            return out
+        else:
+            raise ValueError("Unexpected tensor rank")

The caller should implement the changes necessary to facilitate batched CFG, FWIW. On a separate note, batching CFG for QwenImage Edit didn't lead to any significant gains:

latency=54.8 seconds. (regular)
latency=45.153 seconds. (compile)
latency=54.739 seconds. (batched CFG)
atency=45.024 seconds. (+ compile)

But that is a separate topic and probably shouldn't be discussed here.

"real difference" in my comment: real difference in image quality, if the padding tokens aren't masked.
not real speed difference as you've shown. just clarifying, feel free to mark this comment off-topic.

this would be really great to have resolved, as it lends to a huge disparity in image quality for qwen vs the original impl that did not batch.

@kohya-ss

Here is a fix by @kohya-ss:
kohya-ss/musubi-tuner#688

Might not apply 100% to diffusers repo, because he has copied the code and maybe made other changes before this fix

In my implementation, I swap the order of image and text, and each batch contains only images of the same size (using Aspect Ratio Bucketing), so the implementation is fairly simple.

@sayakpaul

Hi @sayakpaul @naykun,

I'm interested in working on this issue as part of the Diffusers MVP contributor program

Issue Summary:
The encoder_hidden_states_mask parameter is passed to QwenDoubleStreamAttnProcessor2_0 but is never actually used. This breaks batch training when text sequences have different lengths, as padding tokens are not masked during attention computation.

Proposed Fix:
Following the pattern used in the Chroma pipeline, I plan to:

Convert the 1D encoder_hidden_states_mask to a 2D attention mask
Concatenate it with an image attention mask (all ones, since image tokens aren't padded)
Pass the resulting attention mask to dispatch_attention_fn

This is similar to how Chroma handles it here.

I've reviewed the fix in kohya-ss/musubi-tuner#688 which addresses this in their fork. but it seems it was not merged. However, I think there is room to follow diffusers convention more closely:

Uses [text, image] concatenation order (matching transformer_qwenimage.py)
Creates a proper 2D attention mask (query × key) like Chroma does
Works with the existing diffusers attention dispatch system

Does it makes sense? Can I work on it, test and open a PR?

@yiyixuxu

That does make sense to me! @yiyixuxu what do you think?

@naykun please let us know your thoughts, too.

for flash attn you can handle masking by looping over the batch with each individual input being truncated by its mask. it ends up being faster due to the kernels not being swapped out and the result is more accurate than using a block mask. even torch sdpa suffers this limitation of the presence of even a no-op mask disabling the fast path (efficient attention).

@cdutr

Hi @cdutr, thanks for your willingness to contribute to improving Qwen-Image. We leave the attention mask here to support future batched inference, and our main concern is preserving high efficiency across different attention backends. As noted here: #12075 (comment), we originally trained with variable-length attention, but not all backends—notably PyTorch’s native attention—support it. Using an attention mask can be logically correct, but please also consider potential performance impacts when defining default model behavior.

To converge on a final approach, it would be great to add benchmarks across common backends and batch sizes. Happy to keep iterating on this.

cc @dxqb @sayakpaul @yiyixuxu

@naykun

Hi @cdutr, thanks for your willingness to contribute to improving Qwen-Image. We leave the attention mask here to support future batched inference, and our main concern is preserving high efficiency across different attention backends. As noted here: #12075 (comment), we originally trained with variable-length attention, but not all backends—notably PyTorch’s native attention—support it. Using an attention mask can be logically correct, but please also consider potential performance impacts when defining default model behavior.

To converge on a final approach, it would be great to add benchmarks across common backends and batch sizes. Happy to keep iterating on this.

cc @dxqb @sayakpaul @yiyixuxu

@naykun thank you for confirming!
@cdutr please see #12344 (comment) here then for a proposed long-term solution using sequence lengths. As a short-term fix you could still use an attention mask as @naykun described. The current issue is that neither attention mask nor sequence lengths are used.

@naykun

Thanks @naykun for the detailed response. @cdutr let's maybe first try to support this for flash attention and then see how it fares in terms of performance? I think it might be better to focus on an approach that is more long-term in nature.

Thank you all for the fast replies and rich discussion!

Since this is my first contribution to the project, it will be good to start by implementing the attention mask fix with benchmarks focused on flash attention, and measure any performance impact

Once we have those results and if this approach works well, I'd be happy to explore the variable-length sequence approach mentioned in #12344 as a more comprehensive long-term solution

I'll start working on the benchmarks and share the results here soon. Thanks for the guidance!

@cdutr

@cdutr sharing with you my workaround using torch's varlen-fa
I notice that we dont have to call item() with torch api, not sure with flash-attn one @dxqb
with cfg, 2 input prompts, 1 neg prompt, encoder_hidden_states_mask.shape[0] is 1, i assume we can fix it somewhere in the pipeline's forward but i fixed it inside attn processor here

i haven't run benchmarking but output image is expected to be the same with batch size = 1

Details

class QwenDoubleStreamAttnProcessor2_0:
    """
    Attention processor for Qwen double-stream architecture, matching DoubleStreamLayerMegatron logic. This processor
    implements joint attention computation where text and image streams are processed together.
    """
    _attention_backend = None
    _parallel_config = None
    def __init__(self):
        if not hasattr(torch.nn.functional, "scaled_dot_product_attention"):
            raise ImportError(
                "QwenDoubleStreamAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
            )
    def __call__(
        self,
        attn: 'Attention',
        hidden_states: torch.FloatTensor, # Image stream
        encoder_hidden_states: torch.FloatTensor = None, # Text stream
        encoder_hidden_states_mask: torch.FloatTensor = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        image_rotary_emb: Optional[torch.Tensor] = None,
    ) -> torch.FloatTensor:
        if encoder_hidden_states is None:
            raise ValueError("QwenDoubleStreamAttnProcessor2_0 requires encoder_hidden_states (text stream)")
        batch_size = hidden_states.shape[0]
        if encoder_hidden_states_mask.shape[0] != encoder_hidden_states.shape[0]:
            encoder_hidden_states_mask = encoder_hidden_states_mask.repeat(batch_size, 1)
        seq_txt = encoder_hidden_states.shape[1]
        seq_img = hidden_states.shape[1]
        # Compute QKV for image stream (sample projections)
        img_query = attn.to_q(hidden_states)
        img_key = attn.to_k(hidden_states)
        img_value = attn.to_v(hidden_states)
        # Compute QKV for text stream (context projections)
        txt_query = attn.add_q_proj(encoder_hidden_states)
        txt_key = attn.add_k_proj(encoder_hidden_states)
        txt_value = attn.add_v_proj(encoder_hidden_states)
        # Reshape for multi-head attention
        img_query = img_query.unflatten(-1, (attn.heads, -1))
        img_key = img_key.unflatten(-1, (attn.heads, -1))
        img_value = img_value.unflatten(-1, (attn.heads, -1))
        txt_query = txt_query.unflatten(-1, (attn.heads, -1))
        txt_key = txt_key.unflatten(-1, (attn.heads, -1))
        txt_value = txt_value.unflatten(-1, (attn.heads, -1))
        # Apply QK normalization
        if attn.norm_q is not None:
            img_query = attn.norm_q(img_query)
        if attn.norm_k is not None:
            img_key = attn.norm_k(img_key)
        if attn.norm_added_q is not None:
            txt_query = attn.norm_added_q(txt_query)
        if attn.norm_added_k is not None:
            txt_key = attn.norm_added_k(txt_key)
        # Apply RoPE
        if image_rotary_emb is not None:
            img_freqs, txt_freqs = image_rotary_emb
            img_query = apply_rotary_emb_qwen(img_query, img_freqs, use_real=False)
            img_key = apply_rotary_emb_qwen(img_key, img_freqs, use_real=False)
            txt_query = apply_rotary_emb_qwen(txt_query, txt_freqs, use_real=False)
            txt_key = apply_rotary_emb_qwen(txt_key, txt_freqs, use_real=False)
        # build unpad encoder hidden states
        unpad_txt_query = []
        unpad_txt_key = []
        unpad_txt_value = []
        encoder_hidden_states_unpad_lens = encoder_hidden_states_mask.sum(dim=1)
        for original_length, txt_query_i, txt_key_i, txt_value_i in zip(encoder_hidden_states_unpad_lens, txt_query, txt_key, txt_value):
            unpad_txt_query.append(txt_query_i[:original_length])
            unpad_txt_key.append(txt_key_i[:original_length])
            unpad_txt_value.append(txt_value_i[:original_length])

        joint_query = []
        joint_key = []
        joint_value = []
        for i in range(batch_size):
            joint_query.append(torch.cat([unpad_txt_query[i], img_query[i]], dim=0))
            joint_key.append(torch.cat([unpad_txt_key[i], img_key[i]], dim=0))
            joint_value.append(torch.cat([unpad_txt_value[i], img_value[i]], dim=0))
            print(unpad_txt_query[i].shape, img_query[i].shape)
        joint_query = torch.cat(joint_query, dim=0)
        joint_key = torch.cat(joint_key, dim=0)
        joint_value = torch.cat(joint_value, dim=0)
       
        total_lens = encoder_hidden_states_unpad_lens + seq_img
        max_qk = total_lens.max()
        cum_sum_len = torch.cumsum(total_lens, dim=0) - total_lens
        cum_sum_len = cum_sum_len.to(torch.int32)

        # https://gist.github.com/gau-nernst/1d89bb2d064763dbaddbe8cca862adf7
        def varlen_attn(
            query: torch.Tensor,
            key: torch.Tensor,
            value: torch.Tensor,
            cum_seq_q: torch.Tensor,
            cum_seq_k: torch.Tensor,
            max_q: int,
            max_k: int,
            dropout_p: float = 0.0,
            is_causal: bool = False,
        ):
            output, _, _, _, _ = torch.ops.aten._flash_attention_forward(
                query,
                key,
                value,
                cum_seq_q,
                cum_seq_k,
                max_q,
                max_k,
                dropout_p,
                is_causal,
                return_debug_mask=False,
            )
            return output
       
        joint_hidden_states = varlen_attn(
            joint_query,
            joint_key,
            joint_value,
            cum_seq_q=cum_sum_len,
            cum_seq_k=cum_sum_len,
            max_q=max_qk,
            max_k=max_qk,
            dropout_p=0.0,
            is_causal=False,
        )
        # print(joint_hidden_states.shape)
        # Split attention outputs back to padded form, pad zeros to padding length
        heads = joint_hidden_states.shape[1]
        head_dim = joint_hidden_states.shape[2]
        img_attn_output_list = []
        txt_attn_output_list = []
        current_pos = 0
        for i in range(batch_size):
            seg_len = total_lens[i].item()
            seg = joint_hidden_states[current_pos: current_pos + seg_len]
            unpad_len = encoder_hidden_states_unpad_lens[i].item()
            txt_part = seg[:unpad_len]
            img_part = seg[unpad_len:]
            # Pad txt back to seq_txt
            pad_size = seq_txt - unpad_len
            if pad_size > 0:
                padding = torch.zeros((pad_size, heads, head_dim), dtype=seg.dtype, device=seg.device)
                txt_part = torch.cat([txt_part, padding], dim=0)
            img_attn_output_list.append(img_part)
            txt_attn_output_list.append(txt_part)
            current_pos += seg_len
        img_attn_output = torch.stack(img_attn_output_list, dim=0)
        txt_attn_output = torch.stack(txt_attn_output_list, dim=0)
        # Flatten heads into embed_dim
        img_attn_output = img_attn_output.flatten(2)
        txt_attn_output = txt_attn_output.flatten(2)
        # Apply output projections
        img_attn_output = attn.to_out[0](img_attn_output)
        if len(attn.to_out) > 1:
            img_attn_output = attn.to_out[1](img_attn_output) # dropout
        txt_attn_output = attn.to_add_out(txt_attn_output)
        return img_attn_output, txt_attn_output

@bghira

I notice that we dont have to call item() with torch api, not sure with flash-attn one @dxqb with cfg, 2 input prompts, 1 neg prompt, encoder_hidden_states_mask.shape[0] is 1, i assume we can fix it somewhere in the pipeline's forward but i fixed it inside attn processor here

thank you for this, but unfortunately this contains an implicit call to item() in these lines:

        unpad_txt_query.append(txt_query_i[:original_length])
        unpad_txt_key.append(txt_key_i[:original_length])
        unpad_txt_value.append(txt_value_i[:original_length])

this isn't obvious but you'll notice in two ways:

use torch profiler. it'll contain a call to item() and a CUDA synchronisation:

- run it through torch.compile. It'll complain:

torch._dynamo.exc.Unsupported: Dynamic slicing with Tensor arguments
  Explanation: Creating slices with Tensor arguments is not supported. e.g. `l[:x]`, where `x` is a 1-element tensor.
  Hint: It may be possible to write Dynamo tracing rules for this code. Please report an issue to PyTorch if you encounter this graph break often and it is causing performance issues.

  Developer debug context: SliceVariable start: ConstantVariable(NoneType: None), stop: TensorVariable(), step: ConstantVariable(NoneType: None)


from user code:
   File ".../test.py", line 11, in f
    query = txt_query_i[:original_length]

Below is sample code. I don't think there is a way to use flash_varlen or split-attention as described by @bghira above without knowing the sequence lengths on CPU. Reconstructing the sequence lengths from a mask on GPU always requires a transfer back to CPU.
But knowing the sequence lengths on CPU isn't difficult. They just have to be sent to CPU once after the text encoder and then passed along to the attention dispatch. Sending them to CPU in the attention dispatch (or attention itself) causes thousands of CUDA syncs, and is incompatible with torch.compile.

import torch

encoder_hidden_states_mask = torch.tensor([[True, True, True, True, False]], dtype=torch.bool, device='cuda')
txt_query_i = torch.randn((20, 100), device='cuda')

def f():
        encoder_hidden_states_unpad_lens = encoder_hidden_states_mask.sum(dim=1)
        original_length = encoder_hidden_states_unpad_lens[0]
        query = txt_query_i[:original_length]


f()
compiled_f = torch.compile(f, fullgraph=True)
compiled_f()

@dxqb

@dxqb yah absolutely right

@cdutr

@cdutr plan looks good to me. Please let us know once you have something going and we're happy to prioritize it!

@dxqb thanks for the detailed investigation and I concur with your findings as we, ourselves, have done it multiple times to improve torch.compile performance (see 1, 2). We would appreciate your involvement in the corresponding PRs as well.

@dxqb thanks for the detailed investigation and I concur with your findings as we, ourselves, have done it multiple times to improve torch.compile performance (see 1, 2). We would appreciate your involvement in the corresponding PRs as well.

Off-topic: Thanks for this, great read. I think there is more potential in activation quantization though. In my own tests, sampling with activation quantization and linear matmuls in 8 bit is 70% faster than without.

I'm referring to the green and yellow bars, not to the nunchaku one.

I couldn't tell you why since I haven't used torchao, but that it barely helps at all in your link 1, and about 25% faster in your links 2 seems off.

I've implemented a fix for this in PR #12655. The issue was that encoder_hidden_states_mask was accepted but never actually applied in the attention computation.

The fix modifies QwenDoubleStreamAttnProcessor2_0 to create a 2D attention mask from the 1D input mask, properly isolating text padding tokens while keeping all image tokens unmasked. I also adjusted the RoPE positional embeddings to use the full padded sequence length when a mask is present.

I've added tests to verify the mask works correctly and included benchmarks showing ~19% overhead with padding, which is expected for properly handling variable-length sequences in CFG batching. Memory usage is unchanged. Tests on personal computer.

This is my first contribution to diffusers, so feedback is very welcome! I used a custom benchmark (not BenchmarkMixin) to measure the specific impact of mask processing, and named the file benchmarking_qwenimage_mask.py to prevent run_all.py from picking it up since it has a different CSV schema. Happy to adapt if you prefer a different approach.

Looking forward to your thoughts!

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

this should be closed, it was fixed by #12702

Sunbelt Computer Software

PL/B Language Development and Support

Navigation Menu

Repository navigation

[Qwen-image] encoder_hidden_states_mask is not used #12294

Describe the bug

Reproduction

Logs

System Info

Who can help?

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Sunbelt Computer Software

PL/B Language Development and Support

[Qwen-image] encoder_hidden_states_mask is not used #12294

Description

Describe the bug

Reproduction

Logs

System Info

Who can help?

Activity

sayakpaul commented on Sep 6, 2025

dxqb commented on Sep 6, 2025

dxqb commented on Sep 6, 2025

dxqb commented on Sep 7, 2025

sayakpaul commented on Sep 8, 2025

dxqb commented on Sep 8, 2025

sayakpaul commented on Sep 8, 2025

dxqb commented on Sep 8, 2025

bghira commented on Oct 4, 2025

dxqb commented on Oct 26, 2025

kohya-ss commented on Oct 26, 2025

cdutr commented on Nov 11, 2025

sayakpaul commented on Nov 12, 2025

bghira commented on Nov 12, 2025

naykun commented on Nov 12, 2025

dxqb commented on Nov 12, 2025

sayakpaul commented on Nov 12, 2025

cdutr commented on Nov 12, 2025

toilaluan commented on Nov 12, 2025

dxqb commented on Nov 12, 2025

toilaluan commented on Nov 13, 2025

sayakpaul commented on Nov 13, 2025

dxqb commented on Nov 13, 2025

cdutr commented on Nov 14, 2025

github-actions commented on Jan 9, 2026

dxqb commented on Feb 1, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions