add seq_lens to dispatch_attention_fn

Author: kashif

src/diffusers/models/attention_dispatch.py

@@ -305,6 +305,7 @@ def dispatch_attention_fn(
     *,
     backend: Optional[AttentionBackendName] = None,
     parallel_config: Optional["ParallelConfig"] = None,
+    seq_lens: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     attention_kwargs = attention_kwargs or {}
 
@@ -327,6 +328,8 @@ def dispatch_attention_fn(
         **attention_kwargs,
         "_parallel_config": parallel_config,
     }
+    if seq_lens is not None:
+        kwargs["seq_lens"] = seq_lens
     if is_torch_version(">=", "2.5.0"):
         kwargs["enable_gqa"] = enable_gqa
 
@@ -1400,18 +1403,29 @@ def _flash_varlen_attention(
     is_causal: bool = False,
     return_lse: bool = False,
     _parallel_config: Optional["ParallelConfig"] = None,
+    seq_lens: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     batch_size, seq_len_q, _, _ = query.shape
     _, seq_len_kv, _, _ = key.shape
 
-    if attn_mask is not None:
-        attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv)
+    if seq_lens is not None:
+        seq_lens = seq_lens.to(query.device)
+        # use the same lengths for Q and KV
+        seqlens_k = seq_lens
+        cu_seqlens_q = torch.cat([seq_lens.new_zeros(1), seq_lens.cumsum(0)], dim=0).to(torch.int32)
+        cu_seqlens_k = cu_seqlens_q
+        max_seqlen_q = int(seq_lens.max().item())
+        max_seqlen_k = max_seqlen_q
+        attn_mask = None  # varlen uses lengths
+    else:
+        if attn_mask is not None:
+            attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv)
 
-    (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = (
-        _prepare_for_flash_attn_or_sage_varlen(
-            batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device
+        (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = (
+            _prepare_for_flash_attn_or_sage_varlen(
+                batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device
+            )
         )
-    )
 
     key_valid, value_valid = [], []
     for b in range(batch_size):
@@ -1521,18 +1535,28 @@ def _flash_varlen_attention_3(
     is_causal: bool = False,
     return_lse: bool = False,
     _parallel_config: Optional["ParallelConfig"] = None,
+    seq_lens: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     batch_size, seq_len_q, _, _ = query.shape
     _, seq_len_kv, _, _ = key.shape
 
-    if attn_mask is not None:
-        attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv)
+    if seq_lens is not None:
+        seq_lens = seq_lens.to(query.device)
+        seqlens_k = seq_lens
+        cu_seqlens_q = torch.cat([seq_lens.new_zeros(1), seq_lens.cumsum(0)], dim=0).to(torch.int32)
+        cu_seqlens_k = cu_seqlens_q
+        max_seqlen_q = int(seq_lens.max().item())
+        max_seqlen_k = max_seqlen_q
+        attn_mask = None  # varlen uses lengths
+    else:
+        if attn_mask is not None:
+            attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv)
 
-    (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = (
-        _prepare_for_flash_attn_or_sage_varlen(
-            batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device
+        (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = (
+            _prepare_for_flash_attn_or_sage_varlen(
+                batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device
+            )
         )
-    )
 
     key_valid, value_valid = [], []
     for b in range(batch_size):
@@ -2023,21 +2047,31 @@ def _sage_varlen_attention(
     scale: Optional[float] = None,
     return_lse: bool = False,
     _parallel_config: Optional["ParallelConfig"] = None,
+    seq_lens: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     if return_lse:
         raise ValueError("Sage varlen backend does not support setting `return_lse=True`.")
 
     batch_size, seq_len_q, _, _ = query.shape
     _, seq_len_kv, _, _ = key.shape
 
-    if attn_mask is not None:
-        attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv)
+    if seq_lens is not None:
+        seq_lens = seq_lens.to(query.device)
+        seqlens_k = seq_lens
+        cu_seqlens_q = torch.cat([seq_lens.new_zeros(1), seq_lens.cumsum(0)], dim=0).to(torch.int32)
+        cu_seqlens_k = cu_seqlens_q
+        max_seqlen_q = int(seq_lens.max().item())
+        max_seqlen_k = max_seqlen_q
+        attn_mask = None  # varlen uses lengths
+    else:
+        if attn_mask is not None:
+            attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv)
 
-    (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = (
-        _prepare_for_flash_attn_or_sage_varlen(
-            batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device
+        (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = (
+            _prepare_for_flash_attn_or_sage_varlen(
+                batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device
+            )
         )
-    )
 
     key_valid, value_valid = [], []
     for b in range(batch_size):

src/diffusers/models/controlnets/controlnet_qwenimage.py

@@ -228,6 +228,7 @@ def forward(
             joint_attention_kwargs = joint_attention_kwargs.copy()
             lora_scale = joint_attention_kwargs.pop("scale", 1.0)
         else:
+            joint_attention_kwargs = {}
             lora_scale = 1.0
 
         if USE_PEFT_BACKEND:
@@ -246,10 +247,13 @@ def forward(
         temb = self.time_text_embed(timestep, hidden_states)
 
         # Use the encoder_hidden_states sequence length for RoPE computation and normalize mask
-        text_seq_len, encoder_hidden_states_mask = compute_text_seq_len_from_mask(
+        text_seq_len, text_seq_lens_per_sample, encoder_hidden_states_mask = compute_text_seq_len_from_mask(
             encoder_hidden_states, encoder_hidden_states_mask
         )
 
+        if text_seq_lens_per_sample is not None:
+            joint_attention_kwargs.setdefault("text_seq_lens", text_seq_lens_per_sample)
+
         image_rotary_emb = self.pos_embed(img_shapes, text_seq_len, device=hidden_states.device)
 
         timestep = timestep.to(hidden_states.dtype)

src/diffusers/models/transformers/transformer_qwenimage.py

@@ -143,7 +143,7 @@ def apply_rotary_emb_qwen(
 
 def compute_text_seq_len_from_mask(
     encoder_hidden_states: torch.Tensor, encoder_hidden_states_mask: Optional[torch.Tensor]
-) -> Tuple[int, Optional[torch.Tensor]]:
+) -> Tuple[int, Optional[torch.Tensor], Optional[torch.Tensor]]:
     """
     Compute text sequence length without assuming contiguous masks. Returns length for RoPE and a normalized bool mask.
     """
@@ -166,7 +166,7 @@ def compute_text_seq_len_from_mask(
     per_sample_len = torch.where(has_active, active_positions.max(dim=1).values + 1, torch.as_tensor(text_seq_len))
     rope_text_seq_len = max(text_seq_len, int(per_sample_len.max().item()))
 
-    return rope_text_seq_len, encoder_hidden_states_mask
+    return rope_text_seq_len, per_sample_len, encoder_hidden_states_mask
 
 
 class QwenTimestepProjEmbeddings(nn.Module):
@@ -308,6 +308,7 @@ def __call__(
         encoder_hidden_states_mask: torch.FloatTensor = None,
         attention_mask: Optional[torch.FloatTensor] = None,
         image_rotary_emb: Optional[torch.Tensor] = None,
+        text_seq_lens: Optional[torch.Tensor] = None,
     ) -> torch.FloatTensor:
         if encoder_hidden_states is None:
             raise ValueError("QwenDoubleStreamAttnProcessor2_0 requires encoder_hidden_states (text stream)")
@@ -394,6 +395,7 @@ def __call__(
             is_causal=False,
             backend=self._attention_backend,
             parallel_config=self._parallel_config,
+            seq_lens=text_seq_lens,
         )
 
         # Reshape back
@@ -665,6 +667,7 @@ def forward(
             attention_kwargs = attention_kwargs.copy()
             lora_scale = attention_kwargs.pop("scale", 1.0)
         else:
+            attention_kwargs = {}
             lora_scale = 1.0
 
         if USE_PEFT_BACKEND:
@@ -683,10 +686,13 @@ def forward(
         encoder_hidden_states = self.txt_in(encoder_hidden_states)
 
         # Use the encoder_hidden_states sequence length for RoPE computation and normalize mask
-        text_seq_len, encoder_hidden_states_mask = compute_text_seq_len_from_mask(
+        text_seq_len, text_seq_lens_per_sample, encoder_hidden_states_mask = compute_text_seq_len_from_mask(
             encoder_hidden_states, encoder_hidden_states_mask
         )
 
+        if text_seq_lens_per_sample is not None:
+            attention_kwargs.setdefault("text_seq_lens", text_seq_lens_per_sample)
+
         if guidance is not None:
             guidance = guidance.to(hidden_states.dtype) * 1000
 

tests/models/transformers/test_models_transformer_qwenimage.py

@@ -134,9 +134,10 @@ def test_non_contiguous_attention_mask(self):
         encoder_hidden_states_mask[:, 3] = 0
         encoder_hidden_states_mask[:, 5:] = 0
 
-        inferred_rope_len, normalized_mask = compute_text_seq_len_from_mask(
+        inferred_rope_len, per_sample_len, normalized_mask = compute_text_seq_len_from_mask(
             inputs["encoder_hidden_states"], encoder_hidden_states_mask
         )
+        self.assertEqual(int(per_sample_len.max().item()), 5)
         self.assertEqual(inferred_rope_len, inputs["encoder_hidden_states"].shape[1])
         self.assertTrue(normalized_mask.dtype == torch.bool)