Update modeling_quiet.py

modeling_quiet.py

@@ -294,6 +294,7 @@ class QuietAttention(nn.Module):
         self.rope_theta = config.rope_theta
         self.is_causal = True
         self.attention_dropout = config.attention_dropout
+        self._attn_implementation = config._attn_implementation
 
         if (self.head_dim * self.num_heads) != self.hidden_size:
             raise ValueError(
@@ -365,7 +366,30 @@ class QuietAttention(nn.Module):
                 f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
                 f" {attn_weights.size()}"
             )
-
+        if self._attn_implementation == "flash_attention_2":
+            # Prepare attention mask for flash-attn
+            attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
+        elif self._attn_implementation == "sdpa":
+            # Prepare attention mask for SDPA
+            if attention_mask is None or attention_mask.dim() == 2:
+                attention_mask = _prepare_4d_causal_attention_mask(
+                    attention_mask,
+                    (batch_size, seq_length),
+                    inputs_embeds,
+                    past_key_values_length,
+                    sliding_window=self.config.sliding_window,
+                )
+        else:
+            # Prepare attention mask for other implementations
+            if attention_mask is None or attention_mask.dim() == 2:
+                attention_mask = _prepare_4d_causal_attention_mask(
+                    attention_mask,
+                    (batch_size, seq_length),
+                    inputs_embeds,
+                    past_key_values_length,
+                    sliding_window=self.config.sliding_window,
+                )
+                
         if attention_mask is not None:
             if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
                 raise ValueError(
@@ -643,8 +667,63 @@ class QuietFlashAttention2(QuietAttention):
                     causal=causal,
                     window_size=(self.config.sliding_window, self.config.sliding_window),
                 )
+            try:
+                attn_output_unpad = flash_attn_varlen_func(
+                    query_states,
+                    key_states,
+                    value_states,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_k=cu_seqlens_k,
+                    max_seqlen_q=max_seqlen_in_batch_q,
+                    max_seqlen_k=max_seqlen_in_batch_k,
+                    dropout_p=dropout,
+                    softmax_scale=softmax_scale,
+                    causal=causal,
+                )
+            except RuntimeError as e:
+                if "cu_seqlens_q must have shape (batch_size + 1)" in str(e):
+                    # Handle the case when cu_seqlens_q has an invalid shape
+                    if attention_mask is not None:
+                        # Ensure attention_mask has the correct shape
+                        if attention_mask.dim() == 2:
+                            # Convert 2D attention mask to 4D
+                            attention_mask = _prepare_4d_causal_attention_mask(
+                                attention_mask,
+                                (query_states.size(0), query_states.size(1)),
+                                query_states,
+                                past_key_values_length=0,
+                                sliding_window=0,
+                            )
+                        elif attention_mask.dim() != 4:
+                            raise ValueError(
+                                f"Invalid attention mask dimension: {attention_mask.dim()}. Expected 2D or 4D mask."
+                            )
+                        
+                        # Update cu_seqlens_q based on the attention mask
+                        cu_seqlens_q = attention_mask.sum(dim=-1).flatten().cumsum(dim=0).to(torch.int32)
+                        max_seqlen_in_batch_q = cu_seqlens_q[-1].item()
+                        
+                        # Retry flash_attn_varlen_func with updated cu_seqlens_q
+                        attn_output_unpad = flash_attn_varlen_func(
+                            query_states,
+                            key_states,
+                            value_states,
+                            cu_seqlens_q=cu_seqlens_q,
+                            cu_seqlens_k=cu_seqlens_k,
+                            max_seqlen_q=max_seqlen_in_batch_q,
+                            max_seqlen_k=max_seqlen_in_batch_k,
+                            dropout_p=dropout,
+                            softmax_scale=softmax_scale,
+                            causal=causal,
+                        )
+                    else:
+                        raise ValueError(
+                            "Attention mask is required for flash-attn when cu_seqlens_q has an invalid shape."
+                        )
+                else:
+                    raise e
 
-        return attn_output
+            return attn_output
 
     def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
         batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape