Update modeling_quiet.py

modeling_quiet.py

@@ -656,34 +656,47 @@ class QuietFlashAttention2(QuietAttention):
 
         return attn_output
 
-    def upad_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)
+    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
 
-        # Handle the case when seqlens_in_batch is empty
-        if seqlens_in_batch.numel() == 0:
-            max_seqlen_in_batch = 0
-        else:
-            max_seqlen_in_batch = seqlens_in_batch.max().item()
+        # 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)
 
-        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)
+        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)
         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
+            # 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)
+
+        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),
+        )
 # Copied from transformers.models.llama.modeling_llama.LlamaSdpaAttention with Llama->Quiet
 class QuietSdpaAttention(QuietAttention):
     """