Update modeling_quiet.py

modeling_quiet.py

@@ -656,46 +656,34 @@ class QuietFlashAttention2(QuietAttention):
 
         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
-
-        # On the first iteration we need to properly re-create the padding mask
-        # by slicing it on the proper place
-        if kv_seq_len != attention_mask.shape[-1]:
-            attention_mask_num_tokens = attention_mask.shape[-1]
-            attention_mask = attention_mask[:, attention_mask_num_tokens - kv_seq_len :]
-
-        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
-
-        key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
-        value_layer= index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
-        if query_length == kv_seq_len:
-            query_layer = index_first_axis(
-                query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
-            )
-            cu_seqlens_q = cu_seqlens_k
-            max_seqlen_in_batch_q = max_seqlen_in_batch_k
-            indices_q = indices_k
-        elif query_length == 1:
-            max_seqlen_in_batch_q = 1
-            cu_seqlens_q = torch.arange(
-                batch_size + 1, dtype=torch.int32, device=query_layer.device
-            )  # There is a memcpy here, that is very bad.
-            indices_q = cu_seqlens_q[:-1]
-            query_layer = query_layer.squeeze(1)
+    def unpad_input(hidden_states, attention_mask):
+        """
+        Arguments:
+            hidden_states: (batch, seqlen, dim)
+            attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
+        Return:
+            hidden_states: (total_nnz, dim), where total_nnz = number of tokens in selected in attention_mask.
+            indices: (total_nnz,)
+            cu_seqlens: (batch + 1,), use 0 as delimiter.
+            max_seqlen: int
+        """
+        batch_size, seqlen = hidden_states.shape[:2]
+        seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+
+        # Handle the case when seqlens_in_batch is empty
+        if seqlens_in_batch.numel() == 0:
+            max_seqlen_in_batch = 0
         else:
-            # The -q_len: slice assumes left padding.
-            attention_mask = attention_mask[:, -query_length:]
-            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
+            max_seqlen_in_batch = seqlens_in_batch.max().item()
 
-        return (
-            query_layer,
-            key_layer,
-            value_layer,
-            indices_q,
-            (cu_seqlens_q, cu_seqlens_k),
-            (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
-        )
+        indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+        if indices.numel() == 0:
+            indices = torch.zeros(0, dtype=torch.int64, device=hidden_states.device)
+            hidden_states = torch.zeros(0, hidden_states.shape[-1], dtype=hidden_states.dtype, device=hidden_states.device)
+        else:
+            hidden_states = hidden_states.flatten(0, 1)[indices]
+        cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
+        return hidden_states, indices, cu_seqlens, max_seqlen_in_batch
 # Copied from transformers.models.llama.modeling_llama.LlamaSdpaAttention with Llama->Quiet
 class QuietSdpaAttention(QuietAttention):
     """